AtCoder GC023 A: Zero-Sum Ranges

問題

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

解法

最初に問題を見て、これ本当に200点？って思うくらい悩んだ。

とりあえず、累積和をとってみる。

例えば、サイズ4の数列：{a1 a2 a3 a4}の累積和を考える。

S1 = a1 S2 = a1 + a2 S3 = a1 + a2 + a3 S4 = a1 + a2 + a3 + a4

この時、部分数列の和は累積和の差で求められるのがポイント。

例えば、{a3 a4} は、S4 - S2 で求めることが出来る。 つまり、{a3 a4} の和が0になるのは、S4 - S2 = 0、つまり、S4 = S2 な場合といえる。

よって、累積和のうち、等しいものの組み合わせを数え上げることになる。

入力例1：{1 3 -4 2 2 -2}の場合、累積和は{1 4 0 2 4 2}なので、等しいものは、2が2個と4が2個となり、

₂C₂ + ₂C₂

となる。さらに、累積和自体が0の場合はそれ単独でも組み合わせの一つとなり得る。

この場合、累積和の先頭に0を追加してやり、{0 1 4 0 2 4 2}とすることで、同様に扱うことが出来る。

₂C₂ + ₂C₂ + ₂C₂

実装

n = gets.chomp.to_i
a = gets.chomp.split.map(&:to_i)

s = [0]
for i in 0..(n-1)
    s[i+1] = a[i] + s[i]
end

h = Hash.new(0)

s.each do |v|
    h[v] += 1
end

ans = 0
h.each do |k, v|
    next if v == 1
    ans += v * (v-1) / 2
end
puts ans