## 【UVA-11651】Krypton Number System

【UVA-11651】Krypton Number System

Solution:

$$dp[i][k] = \sum dp[i – (j – k)^2][j]$$

#include <iostream>
#include <cstring>

using namespace std;

typedef long long LL;
const LL MOD = 1ll << 32;
const int LIM = 250;
struct Matrix {
LL v[LIM][LIM];
int size;
Matrix operator * (const Matrix & m) const {
Matrix ret = {0};
ret.size = size;
for (int i = 0; i < size; i++) {
for (int k = 0; k < size; k++) {
if (v[i][k]) {
for (int j = 0; j < size; j++) {
ret.v[i][j] = (ret.v[i][j] + v[i][k] * m.v[k][j]) % MOD;
}
}
}
}
return ret;
}
};
int b;
LL s, ans;
LL dp[LIM][LIM];

inline void init() {
memset(dp, 0, sizeof dp);
ans = 0;
int lim = (b - 1) * (b - 1);
for (int j = 1; j < b; j++) dp[0][j] = 1;
for (int i = 0; i <= lim; i++) {
for (int j = 0; j < b; j++) {
for (int k = 0; k < b; k++) {
if (j == k) continue;
int tmp = (j - k) * (j - k);
if (i + tmp > lim) continue;
dp[i + tmp][k] += dp[i][j];
}
}
}
}

Matrix buildBase() {
Matrix ret = {{0}};
ret.size = (b - 1) * (b - 1) * b;
LL tmp;
for (int i = 0; i <= ret.size - b; i++)
ret.v[i][i + b] = 1;
int lim = (b - 1) * (b - 1);
for (int i = 0; i < lim; i++) {
for (int j = 0; j < b; j++) {
for (int k = 0; k < b; k++) {
if (j == k) continue;
tmp = i + (k - j) * (k - j);
if (tmp == lim) ret.v[(tmp - 1) * b + k][i * b + j] = 1;
}
}
}
return ret;
}

Matrix powMod(int pow, Matrix m) {
Matrix ret = {{0}};
ret.size = m.size;
for (int i = 0; i < ret.size; i++) ret.v[i][i] = 1;
while (pow) {
if (pow & 1) {
ret = m * ret;
}
if (pow) m = m * m;
pow >>= 1;
}
return ret;
}

int main(void) {
ios::sync_with_stdio(false);
cin.tie(NULL);
int t;
cin >> t;
for (int k = 1; k <= t; k++) {
cin >> b >> s;
init();
int lim = (b - 1) * (b - 1);
if (s <= lim) {
for (int j = 0; j < b; j++) ans += dp[s][j];
cout << "Case " << k << ": " << ans << "\n";
} else {
Matrix base = buildBase();
int size = lim * b;
int pos = (lim - 1) * b;
base = powMod(s - lim, base);
for (int j = 0; j < b; j++) {
LL tmp = 0;
for (int k = 0; k < size; k++)
tmp = (tmp + base.v[pos + j][k] * dp[k / b + 1][k % b]) % MOD;
ans = (ans + tmp + MOD) % MOD;
}
cout << "Case " << k << ": " << ans << "\n";
}
}
return 0;
}