ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

ABC 133: D問題の解説 (Python3)

問題:AtCoder Beginner Contest 133: D - Rain Flows into Dams

公式解説PDF:https://img.atcoder.jp/abc133/editorial.pdf

公式解説動画:https://youtu.be/8uowVvQ_-Mo?t=3228

解説

考察問題です。数式を変換してXに対する漸化式を導きます。

以下の解説は公式PDFの解説まんまです。

問題の性質により、山に降った雨の総量と、ダムに貯まった水の総量は同じなので、

 A_1+ A_2 + ... A_N = X_1 + X_2 + ... + X_N = S とします。

よって X_1は、

 X_1 = S - (X_2 + ... + X_N)    ・・・(式★)

になります。


また、

 \frac{X_i + X_{i+1}} {2} = A_i

より、

 X_i + X_{i+1} = 2 A_i

なので、(式★)は、

 X_1 = S - 2(A_2 + A_4 + A_{N-1}) と式変形することができます。

これで X_1 は求めることができました。


同様にして X_2, ..., X_N を順に求めて行けばよいですが、単純な2重ループだと O(N^2)になってしまいます。

そこで O(N)で解くために、

 X_i + X_{i+1} = 2 A_i より、

 X_{i+1} = 2 A_i - X_i

の漸化式をつくれるので、それを解いていけばOKです。

コーナーケース

  • 特になし

実装

# -*- coding:utf-8 -*-
import sys


def solve():
    N = int(sys.stdin.readline())
    As = list(map(int, sys.stdin.readline().split()))

    S = sum(As)
    Xs = [0]*N
    # 初期状態を決める(Xs[0])
    Xs[0] = S - 2 * sum(As[1:N:2])

    # 漸化式を解く
    for i in range(N-1):
        Xs[i+1] = 2*As[i] - Xs[i]

    print(*Xs)


if __name__ == "__main__":
    solve()

計算量は O(N) です。