タイトルそのままです。「連続する部分列の総和が0になる場合」が何通りあるのかを数えます。これもループを何重にも回すのではなく1回で終わらせます。累積和をつかえば配列が長くても一発で終わります。
これはAtCoder 200 点問題史上最難問といわれている有名な問題です。
問題:長さ N の整数列 A があります。A の 空でない 連続する 部分列であって、その総和が 0 になるものの個数を求めてください(ただし 1 ≦ N ≦ 100,000、-1,000,000,000 ≦ A[i] ≦ 1,000,000,000)。
連続する部分列の総和が0の場合
累積和 sums を計算したら sums[b] – sums[a](a < b)が0になる部分を探します。これが部分列の総和が0になる部分です。これは言い換えれば「sumsの中から同じ数字を選ぶ組み合わせはいくつあるか?」という問題と同じになります。同じ数が2つしかない場合は「組み合わせは1つ」と数えればよいですが、3つ以上ある場合はどうすればよいでしょうか? n個の中から2つを選ぶ方法は n ×(n - 1) ÷ 2 です。GroupByメソッドをつかってそれぞれ同じ数がいくつあるのかを調べています。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
using System; using System.Collections.Generic; using System.Linq; public class Program { static void Main() { int N = int.Parse(Console.ReadLine()); long[] values = Console.ReadLine().Split().Select(_ => long.Parse(_)).ToArray(); Solve(values); } static void Solve(long[] values) { int len = values.Length; long[] sums = new long[len + 1]; for (int i = 0; i < len; i++) sums[i + 1] += sums[i] + values[i]; var groups = sums.GroupBy(_ => _).ToList(); long ans = 0; foreach (var group in groups) { long count = group.Count(); long key = group.Key; ans += count * (count - 1) / 2; } Console.WriteLine(ans); } } |
連続する部分列の総和が0以外でも使えるか?
ところで「連続する部分列の総和が K になる場合の数を求める」(Kは0とは限らない)場合はどうなるのでしょうか? この場合は辞書にsums[i]を追加し、新しく別のsums[i]を追加するまえにsums[i]から○○を引くとKになるキーが存在するかを調べます。これは言い方を変えるとキーに sums[i] – K が存在するかを調べるということです。
もし見つかったら keyValuePairs[sums[i] – K]を変数 ans に追加します。そして新しいsums[i]も辞書に追加します。すでに存在する場合はkeyValuePairs[value]をインクリメントします。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
public class Program { static void Main() { int N = int.Parse(Console.ReadLine()); long K = int.Parse(Console.ReadLine()); long[] values = Console.ReadLine().Split().Select(_ => long.Parse(_)).ToArray(); Solve(values, K); } static void Solve(long[] values, long K) { int len = values.Length; long[] sums = new long[len + 1]; for (int i = 0; i < len; i++) sums[i + 1] += sums[i] + values[i]; Dictionary<long, long> keyValuePairs = new Dictionary<long, long>(); long ans = 0; foreach (long value in sums) { // value - key = K(書き換えると key = value - K) を満たす key があるか探す if (keyValuePairs.ContainsKey(value - K)) ans += keyValuePairs[value - K]; if (keyValuePairs.ContainsKey(value)) keyValuePairs[value]++; else keyValuePairs.Add(value, 1); } Console.WriteLine(ans); } } |
GroupByメソッドではなくDictionaryを使ったコードに書き換えて提出すると・・・こっちのほうが速い?ようです(引数を変えているので上記コードをそのまま提出してはいけません)。
1回目の誤答はご愛敬ということで・・・。