ベスパリブ

プログラミングを主とした日記・備忘録です。ベスパ持ってないです。

ABC 130: D問題の解説 (Python3)

問題:AtCoder Beginner Contest 130: D - Enough Array

公式解説PDF:https://img.atcoder.jp/abc130/editorial.pdf

公式解説動画:https://youtu.be/ERZuLAxZffQ?t=2473

解説

ABC130のD問題の解説です。しゃくとり法の問題です。

8 10
6 1 2 7 10 1 2 1

という入力を考えます。 この入力は、

  • N = 8
  • K = 10
  • 数列は 6 1 2 7 10 1 2 1

です。

## 初期状態
 6 1 2 7 10 1 2 1
|_|
 6

left = 0  # 部分列の左端の添字(0-origin)
right = 0  # 部分列の右端の添字(0-origin)
ans = 0  # 出力する答え

## rightを右に1進める
 6 1 2 7 10 1 2 1
|___|
  7

## rightを右に1進める
 6 1 2 7 10 1 2 1
|_____|
   9

## rightを右に1進める
 6 1 2 7 10 1 2 1
|_______|
    16

部分列の和がKを超えた。
重要なポイントとして、この時点からrightを右に1移動させた部分列 6 1 2 7 10 の和もKを超える。
さらにrightを右に1移動させた部分列 6 1 2 7 10 1 の和もKを超える。
つまり、この時点から right を一番右に移動させるまでの N-right個 の部分列は必ずKを超えるので、まとめて足してやればよい。
ans += N - right

## total < K になるまで、leftを右に1進める
 6 1 2 7 10 1 2 1
  |_____|
    10

同様に、N-right個をまとめて足してやる。
ans += N - right

## total < K になるまで、leftを右に1進める
 6 1 2 7 10 1 2 1
    |___|
      9

## rightを右に1進める
 6 1 2 7 10 1 2 1
    |______|
       19

ans += N - right

## total < K になるまで、leftを右に1進める
 6 1 2 7 10 1 2 1
      |____|
        17

ans += N - right

## total < K になるまで、leftを右に1進める
 6 1 2 7 10 1 2 1
        |__|
         10

ans += N - right

## total < K になるまで、leftを右に1進める
 6 1 2 7 10 1 2 1
           |
           0

## rightを右に1進める
 6 1 2 7 10 1 2 1
           |_|
            1

## rightを右に1進める
 6 1 2 7 10 1 2 1
           |___|
             3

## rightを右に1進める
 6 1 2 7 10 1 2 1
           |_____|
              4

rightはこれ以上右に進めないので終了

コーナーケース

  • 特になし

実装

# -*- coding:utf-8 -*-
import sys


def solve():
    N, K = list(map(int, input().split()))
    As = list(map(int, sys.stdin.readline().split()))

    left = 0  # 部分列の左端の添字
    total = 0  # 部分列の和
    ans = 0  # 出力する答え
    for right in range(0, N):
        total += As[right]
        while total >= K:
            ans += N - right
            total -= As[left]
            left += 1

    print(ans)


if __name__ == "__main__":
    solve()

計算量は、 O(N)

公式解説によると、二分探索による解法もある。