Codeforces Round #334 (Div. 1) B. Moodular Arithmetic

D 問題にしては簡単すぎない？

問題

codeforces.com

解法

k = 1 の場合のみ特別な場合分けが必要ですが, 基本的な方針は以下のとおりです。

k != 1 のとき, f(0) != 0 であるとすると矛盾するので, f(0) = 0 です。

a を任意の非ゼロ自然数として, a*k^b (b は自然数) という同値類で, 何種類かのグループが出来ます。代表となる数 a における f(a) の数は自由に選ぶことができるので, 答えは p^(グループの数) です。

ですが, k = 1 の場合は, f(0) も好きに選べるので, 答えは p 倍されます。

下のコードではちょっとビビって k = 0 の場合も分けて考えてますが多分必要ないです。

const ll MOD = 1e9+7;

bool done[1000100];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    ll p, k;
    cin >> p >> k;
  if (k == 0) {
      ll ans = 1;
      for (int i = 1; i <= p-1; i++) (ans *= p) %= MOD;
      cout << ans << endl;
  } else if (k == 1) {
      ll ans = 1;
      for (int i = 0; i <= p-1; i++) (ans *= p) %= MOD;
      cout << ans << endl;
  } else {
    {
        ll ans = 1;
        for (int i = 1; i < p; i++) {
            if (!done[i]) {
                ll cur = i;
                while (!done[cur]) {
                    done[cur] = true;
                    cur = (cur*k)%p;
                }
                (ans *= p) %= MOD;
            }
        }
        cout << ans << endl;
    }
    return 0;
}