AtCoder AGC023 A - Zero-Sum Ranges

AtCoder AGC023 A - Zero-Sum Ranges をやってみた。

問題

長さ N の整数列 A があります。
A の 空でない 連続する 部分列であって、その総和が 0 になるものの個数を求めてください。 ただし、ここで数えるのは 部分列の取り出し方 であることに注意してください。 つまり、ある 2 つの部分列が列として同じでも、取り出した位置が異なるならば、それらは別々に数えるものとします。
( https://atcoder.jp/contests/agc023/tasks/agc023_a より引用)

制約

$0 \le N \le 2 \times 10^5$
$-10^9 \le A_i \le 10^9$

考え方

とりあえず 累積和 をとって全パターン数え上げるか、とやってみたら無事 TLE で死亡。（よく考えたら計算量は $O(2^{N+1}-1)$ くらい？）

結局ギブアップで解説を見ました。

累積和の数列 $S$ を作って、$S$ のなかで異なる位置の値が同じ値であれば、その範囲の数値の和が0になることを利用すれば良いようです。
たしかに、累積和の数列を $S$ とすると、普段これを利用して2要素間の合計を計算するには、$S_j - S_i$ などで計算するわけで、これが0になるのは $S_i$ と $S_j$ が同じ値のとき、といえます。

同じ値になる物であれば、Sの要素がいくつかあれば、そこから2つ選べば合計が0になるので、 S の各値について、それぞれ出現数 x とすると、 $2 \le x$ の場合に $ {}_x C_2 $ を計算して合計すれば答えになります。

提出コード

次のようなコードを書きました（抜粋）。全体はこちら

func main() {
    defer stdout.Flush()
    n := readInt()
    a := make([]int64, n)
    for i := range a {
        a[i] = readInt64()
    }
    m := make(map[int64]int64)
    var s int64 = 0
    m[s]++
    for i := range a {
        s = s + a[i]
        m[s]++
    }
    var ans int64
    for _, v := range m {
        if 1 < v {
            ans += v * (v - 1) / 2
        }
    }
    println(ans)
}

余談

drken さんの 累積和の記事 を改めて読んでみたところ、この問題についても解説されていました。