概要

長さ \(N\) の数列 \(A\) から、連続する \(K\) 個の要素を取り出す方法は \(N-K+1\) 通りあります。それらすべてのパターンについて「選んだ要素の総乗（積）」を計算し、その合計を \(10^9 + 7\) で割った余りを求める問題です。

考察

素朴な方法とその限界

各区間について愚直に \(K\) 個の要素を掛け合わせると、1つの区間あたり \(O(K)\) の計算量がかかります。区間の数は \(N-K+1\) 個あるため、全体の計算量は \(O(N \times K)\) となります。 制約では \(N, K \leq 2 \times 10^5\) であるため、最悪の場合 \(4 \times 10^{10}\) 回程度の計算が必要になり、制限時間内に終わりません。

累積積の利用

区間の「和」を高速に求める手法として「累積和」があるように、区間の「積」も累積積（Prefix Product）を利用することで高速化できます。 数列 \(A\) の先頭からの積を保持する配列 \(P\) を用意します。 - \(P_0 = 1\) - \(P_i = A_1 \times A_2 \times \dots \times A_i \pmod{10^9+7}\)

これを用いると、区間 \([i, i+K-1]\) の積は次のように表せます。 $\(\prod_{j=i}^{i+K-1} A_j = P_{i+K-1} \div P_{i-1} \pmod{10^9+7}\)$

逆元の利用

モジュロ演算（余りの計算）において、割り算は「逆元を掛ける」ことと同義です。 \(10^9+7\) は素数であり、制約より \(A_i \geq 1\) なので累積積が \(0\) になることはありません。したがって、フェルマーの小定理を用いることで \(P_{i-1}\) の逆元を求めることができます。

各区間について個別に逆元を計算すると \(O(N \log (\text{mod}))\) かかりますが、累積積の逆元も後ろから順に計算することで \(O(N)\) でまとめて求めることが可能です。

アルゴリズム

1. 累積積の計算: 配列 pre を作成し、pre[i] に \(A\) の先頭 \(i\) 個の積を格納します。 (\(O(N)\))
2. 逆元の計算:
   • まず、全体の積 pre[N] の逆元をフェルマーの小定理（pow(pre[N], mod - 2, mod)）で求めます。
   • 次に、inv_pre[i-1] = inv_pre[i] * A[i-1] という関係を利用して、後ろから順にすべての累積積の逆元を求めます。 (\(O(N)\))
3. 総和の計算: 各開始位置 \(i\) について、pre[i+K] * inv_pre[i] を計算し、それらを合計します。 (\(O(N)\))

計算量

• 時間計算量: \(O(N)\)
  • 累積積の構築に \(O(N)\)、逆元の計算に \(O(N + \log (\text{mod}))\)、最終的な集計に \(O(N)\) かかります。
• 空間計算量: \(O(N)\)
  • 累積積およびその逆元を保持するための配列が必要です。

実装のポイント

• 逆元の一括計算: pow 関数をループ内で何度も呼ぶと重くなるため、最後に一度だけ計算して、あとは順次掛け算で求める手法が効率的です。

• 大きな入力の処理: Pythonでは sys.stdin.read().split() を使うことで、大量の数値を高速に読み込むことができます。

• インデックスの管理: 累積積の配列は、要素数 \(N\) に対してサイズ \(N+1\) で作成すると、空の積（1）を扱いやすくなり実装がスッキリします。

  ソースコード

import sys

def solve():
    # 標準入力からすべてのデータを読み込み、空白で分割してリストにする
    # sys.stdin.read().split() は大量の入力を高速に処理するのに適しています
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # N: 数列の長さ, K: 連続する要素の個数
    n = int(input_data[0])
    k = int(input_data[1])
    
    # A: 数列（0-indexedで管理）
    a = list(map(int, input_data[2:]))
    
    # 法とする素数
    mod = 10**9 + 7
    
    # 累積積（prefix product）を計算する
    # pre[i] は a[0] * a[1] * ... * a[i-1] の積を mod で割った余り
    pre = [1] * (n + 1)
    for i in range(n):
        pre[i+1] = (pre[i] * a[i]) % mod
        
    # 累積積の逆元を O(N) で一括計算する
    # 10^9 + 7 は素数であり、a_i は 1 以上 10^9 以下なので、累積積が 0 になることはない
    inv_pre = [1] * (n + 1)
    # フェルマーの小定理を用いて pre[n] の逆元を計算する
    inv_pre[n] = pow(pre[n], mod - 2, mod)
    # pre[i-1] の逆元は、pre[i] の逆元に a[i-1] を掛けることで求められる
    for i in range(n, 0, -1):
        inv_pre[i-1] = (inv_pre[i] * a[i-1]) % mod
        
    ans = 0
    # 連続する K 個の要素の積の総和を求める
    # 開始インデックス i は 0 から N-K まで
    for i in range(n - k + 1):
        # 連続する K 個の要素 a[i], a[i+1], ..., a[i+k-1] の積は、
        # (a[0]*...*a[i+k-1]) * (a[0]*...*a[i-1])^-1 と等しい
        # すなわち、pre[i+k] * inv_pre[i] で求められる
        current_product = (pre[i+k] * inv_pre[i]) % mod
        ans = (ans + current_product) % mod
        
    # 結果を出力
    print(ans)

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-preview によって生成されました。