しゃくとり法
「いちばん愚直にやると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つ変更を加えてみた。
その1
- これだと尺取り虫の範囲を示すインデックス(l, r)を動かしたときに、範囲内の合計(val)を計算している。
- 累積和を予め計算しておけば、範囲内の合計を計算する必要がない。(累積和[r]-累積和[l] で求まる。)
その2
- 「ある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