C - 連続区間の積の和 / Sum of Products of Consecutive Intervals Editorial by admin
Qwen3-Coder-480B概要
長さ \(N\) の数列から連続する \(K\) 個の要素を選ぶすべての方法について、選んだ要素の積を求め、その総和を \(10^9 + 7\) で割った余りとして出力する問題。
考察
この問題では、連続する \(K\) 個の要素の積をすべて求め、その和を計算する必要があります。素朴な方法としては、各開始位置 \(i\) に対して、\(A_i\) から \(A_{i+K-1}\) までの \(K\) 個の要素を順にかけていく方法があります。しかし、この方法では各積の計算に \(O(K)\) かかり、全体で \(O((N-K+1) \cdot K)\) の計算量となり、最大で \(O(N^2)\) となり、制約 \(N \leq 2 \times 10^5\) では間に合いません(TLE)。
そこで、効率的に連続部分列の積を求める工夫が必要になります。これは累積積を使うことで解決できます。累積積とは、先頭からの積を前もって計算しておくことで、任意の区間の積を高速に求めることができるテクニックです。
例えば、数列 \(A = [a_1, a_2, a_3, a_4]\) の累積積を \(P[i] = a_1 \times a_2 \times \cdots \times a_i\) と定義すると、区間 \([L, R]\) の積は以下のように求められます:
\[ \text{prod}(L, R) = \frac{P[R]}{P[L-1]} \]
ただし、剰余環では除算は通常の意味での除算ではなく、逆元を使って掛け算として表現する必要があります。
したがって、事前に累積積とその逆元を計算しておけば、各区間の積が \(O(1)\) で求められ、全体の計算量が \(O(N)\) になります。
アルゴリズム
- 累積積の前処理
prefix_prod[i]:= \(A_1 \times A_2 \times \cdots \times A_i\) を前計算(mod を取りながら)
- 逆元の累積積の前処理
suffix_inv_prod[i]:= \((A_i \times A_{i+1} \times \cdots \times A_{N})^{-1}\) を後ろから前計算
- 各区間の積の計算
- 各 \(i = 0, 1, ..., N-K\) に対して、区間 \([i, i+K-1]\) の積を以下のように求める: $\( \text{prod} = \text{prefix\_prod}[i+K] \times \text{modinv}(\text{prefix\_prod}[i], MOD) \)$
- この積を答えに足していく(mod を取りながら)
計算量
- 時間計算量: \(O(N)\)
- 空間計算量: \(O(N)\)
実装のポイント
累積積配列は1-basedインデックスで管理すると境界条件がシンプルになる。
逆元の計算は
pow(a, MOD - 2, MOD)を使う(フェルマーの小定理より、\(a^{p-2} \equiv a^{-1} \pmod{p}\))。区間積を求める際に、
prefix_prod[i]が 0 の場合、逆元は存在しないので注意が必要だが、今回は入力がすべて正整数なので問題ない。ソースコード
MOD = 10**9 + 7
def modinv(a, mod):
return pow(a, mod - 2, mod)
def main():
import sys
input = sys.stdin.read
data = input().split()
N = int(data[0])
K = int(data[1])
A = list(map(int, data[2:]))
# 累積積を前処理で求める
# 累積積の添字は1-basedで扱う
prefix_prod = [1] * (N + 1)
for i in range(N):
prefix_prod[i + 1] = (prefix_prod[i] * A[i]) % MOD
# 逆元の累積積も前処理で求める
suffix_inv_prod = [1] * (N + 1)
for i in range(N - 1, -1, -1):
suffix_inv_prod[i] = (suffix_inv_prod[i + 1] * modinv(A[i], MOD)) % MOD
result = 0
# 各連続部分列の積を計算して加算
for i in range(N - K + 1):
# [i, i+K-1] の積 = prefix_prod[i+K] * modinv(prefix_prod[i], MOD)
prod = (prefix_prod[i + K] * modinv(prefix_prod[i], MOD)) % MOD
result = (result + prod) % MOD
print(result)
if __name__ == "__main__":
main()
この解説は qwen3-coder-480b によって生成されました。