公式

C - 連続区間の積の和 / Sum of Products of Consecutive Intervals 解説 by admin

Claude 4.5 Opus

概要

数列から連続する \(K\) 個の要素を選ぶ全ての方法について、それぞれの積を計算し、その総和を求める問題です。スライディングウィンドウとモジュラ逆元を使って効率的に解きます。

考察

素朴なアプローチの問題点

単純に考えると、\(N - K + 1\) 通りの選び方それぞれについて \(K\) 個の要素の積を計算すればよいです。しかし、この方法では: - 選び方の数: \(O(N - K + 1)\) - 各積の計算: \(O(K)\)

合計で \(O((N - K + 1) \times K)\) となり、最悪ケース(\(K \approx N/2\))では \(O(N^2)\) となってTLEになります。

重要な気づき

連続する区間の積を順番に計算する場合、隣り合う区間は \(K - 1\) 個の要素を共有しています。

例えば \(N = 5, K = 3\) の場合: - 区間1: \(A_1 \times A_2 \times A_3\) - 区間2: \(A_2 \times A_3 \times A_4\) - 区間3: \(A_3 \times A_4 \times A_5\)

区間1から区間2に移るとき、\(A_1\) が抜けて \(A_4\) が入るだけです。つまり: $\(\text{区間2の積} = \frac{\text{区間1の積}}{A_1} \times A_4\)$

割り算の問題

\(10^9 + 7\) で割った余りを扱う場合、普通の割り算はできません。代わりにモジュラ逆元を使います。

\(a\) で割ることは、\(a^{-1}\)(\(a\) のモジュラ逆元)を掛けることと同じです。フェルマーの小定理より、\(p\) が素数のとき: $\(a^{-1} \equiv a^{p-2} \pmod{p}\)$

アルゴリズム

  1. 初期化: 最初の \(K\) 個の要素(\(A_1, A_2, \ldots, A_K\))の積を計算し、これを product とする。total に product を加える。

  2. スライディングウィンドウ: \(i = 0, 1, \ldots, N - K - 1\) について以下を繰り返す:

    • \(A_{i+1}\)(0-indexedで \(A[i]\))のモジュラ逆元を計算
    • product を \(A[i]\) で「割り」、\(A[i+K]\) を掛ける
    • 新しい product を total に加える
  3. 結果出力: total を出力

[A0, A1, A2, A3, A4]  (N=5, K=3の例)
 └──────┘            最初の積: A0×A1×A2
     └──────┘        次の積: (A0×A1×A2) ÷ A0 × A3 = A1×A2×A3
         └──────┘    次の積: (A1×A2×A3) ÷ A1 × A4 = A2×A3×A4

計算量

  • 時間計算量: \(O(N \log MOD)\)

    • 初期積の計算: \(O(K)\)
    • スライディングウィンドウ: \(O(N - K)\) 回のループ、各回でモジュラ逆元の計算(繰り返し二乗法)に \(O(\log MOD)\)
    • 合計: \(O(K + (N - K) \log MOD) = O(N \log MOD)\)
  • 空間計算量: \(O(N)\)

    • 配列 \(A\) の保存に \(O(N)\)

実装のポイント

  1. モジュラ逆元の計算: Pythonでは pow(a, MOD - 2, MOD) で \(a^{MOD-2} \mod MOD\) を効率的に計算できます。

  2. オーバーフロー対策: 各計算の後に % MOD を取ることで、値が大きくなりすぎるのを防ぎます。

  3. 0-indexedと1-indexedの混同に注意: コードでは0-indexed、問題文では1-indexedなので、対応を間違えないようにしましょう。

    ソースコード

def solve():
    MOD = 10**9 + 7
    
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # 最初のK個の積を計算
    product = 1
    for i in range(K):
        product = (product * A[i]) % MOD
    
    total = product
    
    # スライディングウィンドウで残りの積を計算
    # A[i]を除いてA[i+K]を追加する
    for i in range(N - K):
        # A[i]で割り、A[i+K]を掛ける
        # 除算はモジュラ逆元を使用
        inv_ai = pow(A[i], MOD - 2, MOD)
        product = (product * inv_ai) % MOD
        product = (product * A[i + K]) % MOD
        total = (total + product) % MOD
    
    print(total)

solve()

この解説は claude4.5opus によって生成されました。