AtCoder ABC 133 D - Rain Flows into Dams - あるエンジニアのAtCoder奮闘記

irisruneです。弊社でRustが布教されているので試していますが、気を付けなければいけないことが多い言語なのでなかなか大変ですね。

問題

単純な連立方程式を解くだけの問題ですが、規模が大きいので一工夫必要となります。

import sys
import numpy as np


def input():
    return sys.stdin.readline().rstrip()


def main():
    n = int(input())
    a = [int(e) for e in input().split()]
    x = [0 for _ in range(n)]
    x[0] = sum(a) - (sum(a[1:n:2]) * 2)
    for i in range(2, n, 2):
        x[i] = x[i-2] + ((a[i-1] - a[i-2]) * 2)
    x[1] = sum(a) - (sum(a[2:n:2]) * 2)
    for i in range(3, n, 2):
        x[i] = x[i-2] + ((a[i-1] - a[i-2]) * 2)
    print(*x)


main()

この問題は $N$ 元連立1次方程式に変換することができます。 $N$ 個の式が作れるため、 $X_1$ は計算量 $O(N)$ で求めることができます。しかし、残りの $X_i$ を同様にして求めると全体の計算量が $O(N^2)$ となってしまいます。

そこで $X_i,X_{i+2}$ の関係について考えると、 $X_{i+2}=X_i+2A_{i+1}-2A_i$ となるため、尺取り法を用いることで $X_3,X_5,\ldots$ を求めることができます。同様に $X_2$ から $X_4,X_6,\ldots$ を求めることができるため、全体の計算量を $O(N)$ に抑えることができます。

$i$ の偶奇にかかわらず漸化式の形は同じなので、上記提出コードのようにループを2つに分ける意味は処理をわかりやすくする以外に特にありません。

雑記

公式解説を見ればわかる通り、 $X_i,X_{i+1}$ についての漸化式が存在します。というより連立方程式のそれぞれの式が漸化式です。なので尺取り法とか考える必要は全くなかったです。
連立方程式を解けさえすればいいので、ソルバーに突っ込んだりすれば非常に簡単に解けるとは思います。
- ただし、行列の形にすると $N\times N$ 行列ができてしまうので、疎行列形式でデータを持つ必要があると思います。
- 実はソルバーについて全く詳しくないのですが、疎行列を渡して計算時間が $O(N^2)$ かからないソルバーってざらにあるものなのでしょうか？