AtCoder Grand Contest 023のA問題を題材に、部分列の総和が0となるような個数の求め方について説明します。

はじめに

ここでは、必要な情報のリンクを載せます。

コンテストへのリンク

累積和の計算をするので、知らない人は下記を参考に確認してください。

まず、AGC023のA問題の解法について簡単に説明します。そして、今回のようなケースにおいて考えるべきポイントのメモを残します。

まず、累積和でN+1の配列Sを下記の様に計算します。

long[] S = new long[N + 1];
for (int i = 0; i < N; i++)
{
    S[i + 1] = S[i] + A[i];
}

ここでS[j] - S[k] (j<k)とすれば、A[j]からA[k-1]までの和を求めることができます。ただ、このすべてのパターンを調べるとなると、O(N²)の計算量がかかります。なので、下記の様に考えます。

S[j]とS[k]が等しいものだけのパターンを考える

等しい値の数から2個選ぶ組み合わせの数を求める

1については、S[j] - S[k]を0にするためには同値である必要があるので、そのパターンだけ考えればよいです。Sをソートすることで数は計算できます。

2については、1で求めた同じ値の数からnC2を求めます。今回はNは2x10⁵までなので通常の公式に当てはめて計算できます。

res += count * (count - 1) / 2; // countは各値の数

問題を解いた時は、累積和の計算後にどう処理したら良いかわからなかったので、今回の問題のように考え方を変えることで計算できる経験を積んで置くと良いですね。