ABC051 B Sum of Three IntegersをO(1)

Last updated Apr 12, 2022

• Editorial
• Fps

今日 maspyさんの形式的べき級数の記事を読んだ後に、サークルの見学会にいったらいきなりバチャが始まって、FPSが使えたのでメモ

ABC051B 全探索で$O(K^2)$ですが、これを$O(1)$まで落としましょう。

まずは、問題をべき級数に落とし込むところからスタートしましょう。$[0, K]$から1つ整数を選ぶのは、$1+x+x^2+…x^K$と表せて、3個選んで和をとるのは$(1+x+x^2+…x^K)^3$となって、求めたい答えは$x^S$の係数となります。この式展開を愚直にすれば$O(K^2)$です。

式変形して計算量を改善していきます。等比数列の和から、$(1+x+x^2+…x^K)^3 = \frac{(1 - x^{K + 1})^3}{(1 - x)^3}$となります。分子は展開して$1 - 3 x^{K + 1} + 3 x^{2 (K + 1)} - x^{3(K + 1)}$となります. $\frac{1}{1 - x} = 1 + x + x^2 + \cdots$は累積和の操作にあたるので、3回累積和をとればいいです. これで$O(K)$まで落ちました。わ〜い！

O(K)提出コード

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

int main() {
  i64 K, S;
  cin >> K >> S;
  K += 1;
  vector<i64> A(3 * K + 1);
  A[0] = 1;
  A[K] = -3;
  A[K * 2] = 3;
  A[K * 3] = -1;
  rep(q,0,3) {
    rep(i,1,3 * K + 1) {
      A[i] += A[i - 1];
    }
  }
  cout << A[S] << endl;
}

3回累積和を取る操作はもっと最適化できそうです。$\frac{1}{(1 - x)^3} = (1 - x)^{-3}$は、負の二項定理を用いて$\sum_{n=0}^{\infty} \binom{n+2}{2} x^n$となるので、これと分子を掛け算すれば良いです。Sの大きさに気をつけて場合分けすれば$O(1)$です！わいわい！

O(1)提出コード

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

int main() {
  i64 K, S;
  cin >> K >> S;
  K += 1;
  i64 ans = 0;
  if(S >= 0) {
    ans += (S + 2) * (S + 1) / 2;
  }
  if(S >= K) {
    i64 n = S - K;
    ans += (-3) * (n + 2) * (n + 1) / 2;
  }
  if(S >= 2 * K) {
    i64 n = S - 2 * K;
    ans += (3) * (n + 2) * (n + 1) / 2;
  }
  if(S >= 3 * K) {
    i64 n = S - 3 * K;
    ans += (-1) * (n + 2) * (n + 1) / 2;
  }
  cout << ans << endl;
}

最後に導いた$O(1)$の解法は、重複組合せを用いた包除原理の解法と一致してます. FPSを使った求め方もおもしろいね〜