プロフィール
AtCoder PR

AGC023AをTypeScriptで解く：累積和とMapで「ゼロ和部分列」をO(N)で数える

 
記事内に商品プロモーションを含む場合があります
Contents
  1. 対象読者
  2. この記事でわかること
  3. 問題概要
  4. ポイント（考え方の芯）
  5. 実装例（TypeScript / Node.js・1パスO(N)）
  6. よくある間違い
  7. 実装例（発展）
  8. まとめ

対象読者

AtCoder初級〜中級。**「連続部分列の和が0になる個数を効率よく数えたい」**人向け。

この記事でわかること

  • ゼロ和部分列は**「同じ累積和が出現した回数の組み合わせ」**で数えられる
  • 1パス・O(N)・追加配列不要の実装（TypeScript/Node.js）
  • 初心者がハマりやすい落とし穴と回避方法
  • 発展的な別解（ソートで数える／BigInt完全対応）

問題概要

  • 長さNの整数列Aが与えられる
  • 空でない連続部分列の総和が0になるものの個数を数える
  • 制約：N ≤ 2×10^5, |A_i| ≤ 10^9

直感的には「全部の区間を試す」二重ループを思いつきがちですが、N=2e5ではO(N^2)で間に合いません。コツは“累積和（prefix sum）”を使って区間和を引き算にすることです。

ポイント（考え方の芯）

結論：答えは「同じ累積和が現れるインデックスの組(i<j)の個数」です。

理由：
先頭からの累積和をS[0]=0, S[k]=A+…+A[k]とします。
区間[i+1..j]の和は S[j]−S[i]。これが0になるのは S[j]=S[i] のとき。
つまり、同じ値の累積和が出現した回数vごとに C(v,2) を足せば総数になります。
さらに1パスで数えるなら、走査中に「その累積和がすでに何回出たか」を数え、到達時にその回数だけ答えに加えると、C(v,2)の総和が自然に積み上がります。

例（入力例1）：A =
累積和 S: 0, 1, 4, 0, 2, 4, 2
出現回数：

  • 0が2回 → C(2,2)=1
  • 4が2回 → 1
  • 2が2回 → 1
    合計3つ（(1,3,-4), (-4,2,2), (2,-2)）。

もう一度結論：「初期値S=0を1回カウントしておく」「累積和ごとの出現回数をMapで持つ」「到達時に現在までの出現回数を答えに加える」—この3点が鍵です。

実装例（TypeScript / Node.js・1パスO(N)）

  • 計算はnumberで十分ですが、答えの個数は最大で約2e10に届くためBigIntで安全に加算します（出力時にtoString()）。
import * as fs from "fs";

const tokens = fs.readFileSync(0, "utf8").trim().split(/\s+/);
const N = Number(tokens[0]);
const A = tokens.slice(1).map(Number);

let sum = 0;                          // 累積和（numberでOK）
const freq = new Map<number, number>();
freq.set(0, 1);                       // S=0 を1回カウント（先頭からの区間用）
let ans = 0n;                         // 個数はBigIntで安全に

for (let i = 0; i < N; i++) {
  sum += A[i];
  const c = freq.get(sum) ?? 0;       // これまでに同じ累積和が何回出たか
  ans += BigInt(c);                   // その回数分、ゼロ和区間が増える
  freq.set(sum, c + 1);               // 今回のsumを記録
}

console.log(ans.toString());

計算量とメモリ

  • 時間計算量：O(N)
  • 追加メモリ：累積和の種類数に比例（最悪でもO(N)）

よくある間違い

  1. S=0を初期カウントし忘れる
    先頭からの区間（）が取りこぼされます。必ず freq.set(0,1)。
  2. 二重ループで全区間を試す（TLE）
    N=2e5では間に合いません。累積和＋MapでO(N)に。
  3. Objectで頻度表を作る
    "-1"と1などのキー衝突や、プロトタイプ影響を避けるためMap推奨。
  4. BigIntとnumberを混算
    JSでは混ぜられません。上の実装のようにsumはnumber、ansだけBigIntにするか、次の「発展」版のように全てBigIntに統一。
  5. 出力でBigIntのままconsole.log(ans)
    そのままでも出ますが、AtCoderの採点や環境によっては明示が安心。**toString()**を付けましょう。

実装例（発展）

1) 「完全安全」— 累積和も頻度もBigIntで統一

import * as fs from "fs";

const t = fs.readFileSync(0, "utf8").trim().split(/\s+/);
const N = Number(t[0]);
const A = t.slice(1).map(BigInt);

let sum = 0n;
const freq = new Map<bigint, bigint>();
freq.set(0n, 1n);
let ans = 0n;

for (let i = 0; i < N; i++) {
  sum += A[i];
  const c = freq.get(sum) ?? 0n;
  ans += c;
  freq.set(sum, c + 1n);
}

console.log(ans.toString());

  • 多少遅くなりますが、境界値も気にせず書けます。

2) 別観点で理解する—「累積和を並べて、同値の塊ごとにC(v,2)」

  1. S[0..N]を配列に格納してソート
  2. 連続する同値の塊のサイズをvとして**v*(v-1)/2**を足す
  • 計算量はO(N log N)。理屈を掴むのに◎、実戦は1パス解が推奨。

まとめ

  • ゼロ和区間 ⇔ 同じ累積和のペア
  • 実装の肝は freq.set(0,1) と 到達時加算（ans += freq.get(sum)）
  • O(N)/1パス/MapでAtCoderの制約を余裕で満たせます。
  • 慣れたらBigInt統一やソート別解で理解を深めましょう

RELATED POST