AtCoder Grand Contest 023: A - Zero-Sum Ranges

競技プログラミング AtCoder 累積和

問題

問題文

https://atcoder.jp/contests/agc023/tasks/agc023_a

問題概要

長さ N の整数列 A が与えられる.

A の空でない連続する部分列のうち, 総和が 0 となるものの個数を求めよ.

制約

  • 1 \leqq N \leqq 2 \times 10^{5}

  • -10^{9} \leqq A_i \leqq 10^{9}

解答例

指針

  • 連続する部分列の総和 => 累積和

解説

初項を 0 とし, A の累積和をとった数列を S とする.

このとき, A の区間 [l, r] 要素の和は S_{r+1} - S_l として求めることができる.

例えば A = {1, 3, -4, 2, 2, -2} であるとき,

累積和をとった数列は S = {0, 1, 4, 0, 2, 4, 2} となり,

A[1] + A[2] + A[3] = S[4] - S[1] = 1 (0-based-index で考えた)

ここで連続する部分列の総和が 0 になるとは

A[l] + A[l+1] + ... A[r] = S[r+1] - S[l] = 0

<=> S[r+1] = S[l]

すなわち, S[r+1] = S[l] (ただし l < r+1) と同値である.

数列 S の要素を見ていき等しい値となるものが現れたら, それまでに現れているその値の要素数だけ答えとなるの値を増やすというような操作をすることで A の部分列のうち総和が 0 となる区間の数を求めることができる.

等しい値となる要素の数え方としては sort して隣り合う者同士を比較する方法と連想配列を使う方法がある.

  • C++ による実装例
// submission: https://atcoder.jp/contests/agc023/submissions/4650948
// Language: C++14 (GCC 5.4.1)
// Time: 146 ms
// Memoly: 12800 KB

#include <iostream>
#include <map>
 
using namespace std;
 
typedef long long ll;
 
int main() {
    int N; cin >> N;
    map<ll, int> mp;
 
    mp[0] = 1;
 
    ll s = 0, ans = 0;
 
    for (int i = 0; i < N; i++) {
        ll a; cin >> a;
        s += a;
        ans += mp[s];
        mp[s]++;
    }
 
    cout << ans << endl;
 
    return 0;
}
  • Python3 による実装例
# submission: https://atcoder.jp/contests/agc023/submissions/4651016
# Language: Python3 (3.4.3)
# Time: 194 ms
# Memoly: 45788 KB

from collections import defaultdict

N = int(input())
A = map(int, input().split())

d = defaultdict(int)
d[0] = 1

s, ans = 0, 0

for a in A:
    s += a
    ans += d[s]
    d[s] += 1

print (ans)

参考文献