情報自炊

情報を血肉にしたい

【Ruby】AtCoder Beginner Contest 130のD問題振り返り

問題

atcoder.jp

答え

n, k = gets.chomp.split(" ").map(&:to_i)
a_arr = gets.chomp.split(" ").map(&:to_i)

count = 0
j = 0
sum = 0
# iを一つずつ右に進めるようなループ
(0..n-1).each do |i|
  # jを一つずつ右進めるようなループ
  while sum < k && j < n do
    sum += a_arr[j]
    j += 1 # jは進むだけで後戻りしない
  end
  # jを最後まで進めても条件を満たさない場合,
  # 以降iを進めてもsumが小さくなるだけなのでbreak
  break if sum < k
  # a_arr[i..j].sum >= k ということは
  # a_arr[i..j+1].sum, a_arr[i..j+2].sum, ... もk以上なので
  # まとめてカウントアップ
  count += n - j + 1
  sum -= a_arr[i]
end

puts count

jを0から(あるいはiから)始まるループを作ってしまうと最悪実行時間がO(n)になってしまうので、全体でO(n^2)になってしまう。jを戻さないのがポイント。

jを戻さなくて良い理由はよく考えると分かる。

a_arr[i..j]で初めてk以上になったとき、次のループはa_arr[(i+1)..j]から始まる。jを好きなだけ戻した部分列は正の整数mを用いてa_arr[(i+1)..(j-m)]のように表すことができる。

このとき a_arr[(i+1)..(j-m)] >= k が成り立つとすると、a_arr[(i+1)..(j-m)]a_arr[i..(j-m)]の部分列なのでa_arr[i..(j-m)] > a_arr[(i+1)..(j-m)] >= k が成り立つ。これは「a_arr[i..j]で初めてk以上になった」という点と矛盾する。だから常にa_arr[(i+1)..(j-m)] < kとなるのでjを戻すのは考えなくて良い。

このアルゴリズムはしゃくとり法と呼ばれるもので私は初めて知った。

デバッグコードあり

n, k = gets.chomp.split(" ").map(&:to_i)
a_arr = gets.chomp.split(" ").map(&:to_i)

count = 0
j = 0
sum = 0
(0..n-1).each do |i|
  puts "-------- i=#{i} j=#{j}"
  while sum < k && j < n do
    print "add #{a_arr[j]} (j=#{j}). sum: #{sum}->"
    sum += a_arr[j]
    puts "#{sum}"
    j += 1
  end
  puts "after while. sum=#{sum} j=#{j}"
  break if sum < k
  print "add #{n} - #{j} + 1. count: #{count}->"
  count += n - j + 1
  puts "#{count}"
  sum -= a_arr[i]
  puts "diff #{a_arr[i]} (i=#{i}). sum=#{sum}"
end

puts count

せっかくなのでデバッグコードを仕込みまくって処理を追った。こうするとインデックスの動きがよく見えて分かりやすい。

入力例1

4 10
6 1 2 7
-------- i=0 j=0
add 6 (j=0). sum: 0->6
add 1 (j=1). sum: 6->7
add 2 (j=2). sum: 7->9
add 7 (j=3). sum: 9->16
after while. sum=16 j=4
add 4 - 4 + 1. count: 0->1
diff 6 (i=0). sum=10
-------- i=1 j=4
after while. sum=10 j=4
add 4 - 4 + 1. count: 1->2
diff 1 (i=1). sum=9
-------- i=2 j=4
after while. sum=9 j=4
2

入力例2

3 5
3 3 3
-------- i=0 j=0
add 3 (j=0). sum: 0->3
add 3 (j=1). sum: 3->6
after while. sum=6 j=2
add 3 - 2 + 1. count: 0->2
diff 3 (i=0). sum=3
-------- i=1 j=2
add 3 (j=2). sum: 3->6
after while. sum=6 j=3
add 3 - 3 + 1. count: 2->3
diff 3 (i=1). sum=3
-------- i=2 j=3
after while. sum=3 j=3
3

入力例3

10 53462
103 35322 232 342 21099 90000 18843 9010 35221 19352
-------- i=0 j=0
add 103 (j=0). sum: 0->103
add 35322 (j=1). sum: 103->35425
add 232 (j=2). sum: 35425->35657
add 342 (j=3). sum: 35657->35999
add 21099 (j=4). sum: 35999->57098
after while. sum=57098 j=5
add 10 - 5 + 1. count: 0->6
diff 103 (i=0). sum=56995
-------- i=1 j=5
after while. sum=56995 j=5
add 10 - 5 + 1. count: 6->12
diff 35322 (i=1). sum=21673
-------- i=2 j=5
add 90000 (j=5). sum: 21673->111673
after while. sum=111673 j=6
add 10 - 6 + 1. count: 12->17
diff 232 (i=2). sum=111441
-------- i=3 j=6
after while. sum=111441 j=6
add 10 - 6 + 1. count: 17->22
diff 342 (i=3). sum=111099
-------- i=4 j=6
after while. sum=111099 j=6
add 10 - 6 + 1. count: 22->27
diff 21099 (i=4). sum=90000
-------- i=5 j=6
after while. sum=90000 j=6
add 10 - 6 + 1. count: 27->32
diff 90000 (i=5). sum=0
-------- i=6 j=6
add 18843 (j=6). sum: 0->18843
add 9010 (j=7). sum: 18843->27853
add 35221 (j=8). sum: 27853->63074
after while. sum=63074 j=9
add 10 - 9 + 1. count: 32->34
diff 18843 (i=6). sum=44231
-------- i=7 j=9
add 19352 (j=9). sum: 44231->63583
after while. sum=63583 j=10
add 10 - 10 + 1. count: 34->35
diff 9010 (i=7). sum=54573
-------- i=8 j=10
after while. sum=54573 j=10
add 10 - 10 + 1. count: 35->36
diff 35221 (i=8). sum=19352
-------- i=9 j=10
after while. sum=19352 j=10
36