しゃくとり法

2019-06-19

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

しゃくとり法の存在はなんとなく知っていたが、そうかこういう場合に使うと一気に 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