Ruby で解く AtCoder ABC 233 D

はじめに

AtCoder Problems の Recommendation を利用して、過去の問題を解いています。
AtCoder さん、AtCoder Problems さん、ありがとうございます。

今回のお題

Difficulty: 726

解法

今回のテーマ：累積和の応用

入力例 1 
6 5
8 -3 5 7 0 -4

まず、累積和を取ります

入力	累積和
	0
8	8
-3	5
5	10
7	17
0	17
-4	13
大人の都合で0が入っていますが、通常の累積和とは異なり、単調増加になっていないところが応用編となります。
累積和をグラフにしましたが、横軸0の地点から左に山登りをすると考えます。
今の地点から相対的な高さk（入力例１では5）だけ高い地点を見つけます。
左から右に進んで、合計３つのkだけ高い地点を見つけました。

実装例1

n, k = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
r = Array.new(n + 1, 0)
h = { 0 => [0] }
1.upto(n) do |i|
  r[i] = r[i - 1] + a[i - 1]
  if h[r[i]]
    h[r[i]] << i
  else
    h[r[i]] = [i]
  end
end
cnt = 0
0.upto(n) do |i|
  if h[r[i] + k] && h[r[i] + k].bsearch_index{ _1 > i }
    cnt += h[r[i] + k].size - h[r[i] + k].bsearch_index{ _1 > i }
  end
end
puts cnt

TLE対策として、連想配列hを使用します。
また、左から右に進みますので、hの中身は単調増加な配列となり、二分探索bsearchが利用可能になります。

実装例2［追記］

n, k = gets.split.map(&:to_i)
a = gets.split.map(&:to_i)
r = [0]
h = Hash.new(0)
h[0] = 1
cnt = 0
1.upto(n) do |i|
  r[i] = r[i - 1] + a[i - 1]
  cnt += h[r[i] - k]
  h[r[i]] += 1
end
puts cnt

いつも参考にさせていただいている @u2dayo さんのコードを移植したものです。
実装例1はr[i] + kでしたが、こちらはr[i] - kで、山登りで振り返っているイメージです。
この場合、二分探索を省略できますので、実装が軽くなり素晴らしいです。

まとめ

• ABC 233 D を解いた
• 累積和 に詳しくなった