Official

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)\) になります。

アルゴリズム

  1. 累積積の前処理
    • prefix_prod[i] := \(A_1 \times A_2 \times \cdots \times A_i\) を前計算(mod を取りながら)
  2. 逆元の累積積の前処理
    • suffix_inv_prod[i] := \((A_i \times A_{i+1} \times \cdots \times A_{N})^{-1}\) を後ろから前計算
  3. 各区間の積の計算
    • 各 \(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 によって生成されました。