ATABC293E Geometric Progression

ATABC293E Geometric Progression

感觉自己写了一个神秘做法。

题意

给你三个整数 \(A\)，\(M\) 与 \(X\)，求：

\[\sum_{i = 0}^{X - 1} A^i \]

对 \(M\) 取模后的值。

\(1 \le A, m \le 10^9\)，\(1 \le X \le 10^{12}\)。

思路

考场上第一眼看这个题想到的是用这个式子：

\[\sum_{i = 0}^{X - 1} A^i = \frac{A^{X} - 1}{A - 1} \]

然后除法用逆元做一下。结果忽略了一个很重要的事实，就是题目没有保证 \(A - 1\) 与 \(M\) 互质，连续 WA 了 3 次才发觉不对劲。

这道题的官方题解是矩阵快速幂？我偏不用其实是不会！

下面是我做这道题的思路。

由于我们有 \(\sqrt{10^{12}} = 10^6\)，所以我们可以从这方面找找思路（谁会这么想啊喂），考虑将原式分成 \(\sqrt{X}\) 块后再做前缀和，设 \(B = \lfloor\sqrt{X}\rfloor\)，记第一块的值为 \(sum\)，有：

\[sum= \sum_{i = 0}^{B - 1} A^i \]

然后通过一个小小的递推得到每一块的值。具体地，有第 \(i\) 块的值为 \(sum \times A^{(i - 1)B}\)。最后剩一个散块暴力累加一下就好了。

那么最终的答案就为：

\[ans = \sum_{i = 1}^{B} sum \times A^{(i - 1)B} + \sum_{i = B^2}^{X - 1} A^i \]

由于我太傻了，是直接在循环中用快速幂计算幂次的，所以我的时间复杂度为 \(O(\sqrt{X}\log{X})\) 级别的，但事实上我们可以预处理出 \(\sqrt{X}\) 内的 \(A\) 的幂次，这样时间复杂度为 \(O(\sqrt{X})\)。虽然不如其他的神仙做法，不过已经足够好了。

代码

代码其实非常好写，毕竟是优雅的暴力。

我在一些奇怪的地方用了 __int128，不过事实上完全不用，只是我赛时发癫。

#include <bits/stdc++.h>

using i64 = long long;
using i128 = __int128;

i64 fastPow(i64 base, i64 pow, i64 mod) {
  i64 res = 1;
  while (pow) {
    if (pow & 1) {
      (res *= base) %= mod;
    }
    (base *= base) %= mod;
    pow >>= 1;
  }
  return res;
}

int main() {
  std::ios::sync_with_stdio(false);
  std::cin.tie(0);

  i64 a, m, x;
  std::cin >> a >> x >> m;

  i64 powNum = 0;
  int block = std::sqrt(x);
  for (int i = 0; i < block; i++) {
    (powNum += fastPow(a, i, m)) %= m;
  }

  i64 ans = 0;
  for (int i = 1; i <= block; i++) {
    (ans += fastPow(a, ((i64) (i - 1) * block), m) * i128(powNum)) %= m;
  }

  if ((i64) block * block != x) {
    for (i64 i = (i64) block * block; i < x; i++) {
      (ans += fastPow(a, i, m)) %= m;
    }
  }

  std::cout << ans << "\n";

  return 0;
}