ABC267 C - Index × A(Continuous ver.) Python解答例

AtCoder Beginner Contest 267 C - Index × A(Continuous ver.)をPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

長さ$N$の整数列$A=(A_1,A_2,\dots,A_N)$が与えられます。

長さ$M$の$A$の連続部分列$B=(B_1,B_2,\dots,B_M)$に対する、$\displaystyle \sum_{i=1}^{M} i \times B_i$の最大値を求めてください。

制約

• $1 \le M \le N \le 2 \times 10^5$
• $- 2 \times 10^5 \le A_i \le 2 \times 10^5$
• 入力は全て整数。

解答例

数式変形と累積和

愚直に解くと$\mathcal{O}(MN)$かかってしまうので、効率よく解く方法を考えます。

以下、プログラムを実装するときに混乱しないように、配列を0-indexedで表します。
すなわち、$A = (A_0, A_1, \ldots, A_{N - 1})$、$B = (B_0, B_1, \ldots, B_{M - 1})$に対して、$\sum^{M - 1}_{i = 0} i \times B_i$を考えます。

とりあえず数式を変形することを考えます。整数列$B$の定義を、ある整数$k(0 \leq k \leq N - M)$を用いて、$B = (A_k, A_{k + 1}, \ldots, A_{M + k - 1})$とします。

これを用いると、求める解の式は以下のように変形することができます。

$$\begin{align*} \sum^{M - 1}_{i = 0} i \times B_i &= \sum^{M + k - 1}_{i = k} (i - k) \times A_{i} \\ &= \sum^{M + k - 1}_{i = k} (i \times A_{i} - k \times A_{i}) \\ &= \sum^{M + k - 1}_{i = k} i \times A_{i} - k \times \sum^{M + k - 1}_{i = k} A_{i} \\ \end{align*}$$

この式変形により、$\Sigma$の中身が$k$に依存しなくなりました。従って、以下の2つの累積和を事前に計算しておくことで区間和を高速に計算できることになります。

$$\begin{align*} C_j &= \begin{cases} 0 \quad (j = 0) \\ \sum^{j - 1}_{k = 0} A_k \quad (1 \leq j \leq N) \end{cases} \\ D_j &= \begin{cases} 0 \quad (j = 0) \\ \sum^{j - 1}_{k = 0} k \times A_k \quad (1 \leq j \leq N) \end{cases} \\ \end{align*}$$

累積和の配列$C, D$は長さが$N + 1$になることに注意してください。^[1]

累積和を用いると、$\sum^{M + k - 1}_{i = k} A_i = C_{M + k} - C_k$、$\sum^{M + k - 1}_{i = k} i \times A_i = D_{M + k} - D_k$となるので、求める解は

$$\begin{align*} \sum^{M - 1}_{i = 0} i \times B_i &= \sum^{M + k - 1}_{i = k} (i - k) \times A_{i} \\ &= \sum^{M + k - 1}_{i = k} (i \times A_{i} - k \times A_{i}) \\ &= \sum^{M + k - 1}_{i = k} i \times A_{i} - k \times \sum^{M + k - 1}_{i = k} A_{i} \\ &= (D_{M + k} - D_{k}) - k \times (C_{M + k} - C{k}) \end{align*}$$

となります。全ての$k$について解を求めたときの最大値が答えとなります。

実装例

Pythonによる実装例を以下に示します。累積和はitertools.accumulateを用いることで1行で記述できます。

c.py

from typing import *
import itertools


def main():
    # 入力受け取り
    N, M = map(int, input().split())
    A: List[int] = list(map(int, input().split()))

    # 累積和配列の作成
    # 右半開区間[l, r)を扱うために先頭に0を付け加えている
    C: List[int] = [0] + list(itertools.accumulate(A))
    D: List[int] = [0] + list(
        itertools.accumulate([(i + 1) * a for i, a in enumerate(A)])
    )

    # 整数kを全探索
    # 今回の問題では負の数を扱うので、解の初期値には十分小さい負の値を入れておく
    answer: int = -(1 << 60)
    for k in range(N - M + 1):
        X: int = D[k + M] - D[k]
        Y: int = k * (C[k + M] - C[k])
        answer = max(answer, X - Y)

    # 答えの出力
    print(answer)

if __name__ == "__main__":
    main()

実際の提出はこちら。

脚注

1. 長さ$N$の累積和配列$C$を作ったとき、整数$l, r(0 \leq l < r < N)$に対して$C_r - C_l$は左半開区間$(l, r]$の区間和となります。競プロの文脈だと左半開区間を扱うことはそんなに無く、むしろ右半開区間$[l, r)$に関心があることがほとんどです。右半開区間$[l, r)$の区間和を求めたいときは$C_{r - 1} - C_{l - 1}$とすればよいですが、添字ミスでバグを生みやすいので良い手法ではないです。
   そこで、先頭に0を付け加えた長さ$N + 1$の累積和配列$C'$を用いると、上と同じ整数$l, r$について右半開区間$[l, r)$の区間和を$C'_{r} - C'_{l}$で求めることができるようになります。こちらの方が簡潔に記述できるようになります。(どのみち添字の扱いには注意しなければなりませんが)