しゃくとり法

[https://atcoder.jp/contests/abc130/tasks/abc130_d:title]

「いちばん愚直にやるとO(n**3), 累積和つかってもO(n**2)だな」まで考えたところで詰まったので解説みた。

しゃくとり法の存在はなんとなく知っていたが、そうかこういう場合に使うと一気にO(n)にできるんだ。

しゃくとり法のアイデアだけおさえて書いたらACできた。コードは以下。

n, k = gets.split.map(&:to_i)
aa = gets.split.map(&:to_i)

val = aa[0]
ans = 0
r = 0

(0...n).each do |l|
    val -= aa[l-1] if l > 0
        while r < n - 1 do
        (ans += n-r) && break if val >= k
        r += 1
        val += aa[r]
    end
    if r == n-1
    ans += 1 if val >= k
    end
end

p ans

ここから2つ変更を加えてみた。

• その１

  • これだと尺取り虫の範囲を示すインデックス(l, r)を動かしたときに、範囲内の合計(val)を計算している。
  • 累積和を予め計算しておけば、範囲内の合計を計算する必要がない。（累積和[r]-累積和[l] で求まる。）

• その２

  • 「あるlに対してrを右に動かして、kを超えたところより右側に長い部分列を全てansに加える」と考えたが、これだとrが一番右にいったときに「lを右に動かしてkを超えたらその部分列（のみを）ansに加える」必要があった。（その部分列「のみを」加えるのは、「l-1番目からrまで」とか「l-2番目からrまで」とかいう部分列は既に数えているから。）これだと2つの数え上げ方を組み合わせるのが気持ちよくない。
  • 「あるrに対してlを(rを超えない範囲で)右に動かして、kを下回る直前の部分列よりも左側に長い部分列を全てansに加える」とすると、数え上げ方が1つで済んで気持ちいい。

変更を加えた後のコードが以下。

n, k = gets.split.map(&:to_i)
aa = gets.split.map(&:to_i)

cs = 0
acs = [0]
aa.each{|a| acs << acs[-1] + a}

l = 0
ans = 0

(1..n).each do |r|
  while l < r do
    break if acs[r] - acs[l] < k
    l += 1
  end
  ans += l
end

p ans