初心者のプログラミング日記

プログラミング初心者の日記

プログラミングに関することを書いていきます。

しゃくとり法

しゃくとり法に関しては下記URLのスライド55ページ~を参考にどうぞ
https://www.slideshare.net/yuki2006_debel/21-35882983

ポイントとしては条件に一致するまでrightを動かしていき、処理が終わったらleftを動かす。この時rightの値はそのままにする。

問題は以下を使います
https://atcoder.jp/contests/abc130/tasks/abc130_d

入力まで

N,K=map(int,input().split())
a=list(map(int,input().split()))
right=0 #どこまで探索したか
ans=0 #答え
num=0 #left~rightの合計値

処理部分を書いていきます

for left in range(N):
    while num<K and right<N: #numがKを超えるかrightがNを超えたら終了
        num+=a[right] #値を足す
        right+=1 #rightを動かす

    if num>=K: #numがk以上
        ans+=N-right+1#left=rightが条件を満たしているため、rightに+1しても条件を満たす
        num-=a[left] #leftを動かすためnumからa[left]を引く

しゃくとり法を使うことによってO(n)で解くことができます。
while文の条件は今回の問題の場合は、連続部分列に含まれる全ての要素の値の和はK以上である必要があるので一つ目の条件はそれを満たすようにします。
もう1つの条件はrightは端までいったらループを終了します。
if文はleftを動かている途中でnumがK以下になる場合があるのでそれを弾いています。
ansの所は例えば例題2の場合[3,3,3]という配列が与えられ、left=0,right=2の場合[3,3],[3]に区間が分けられる。その場合左の区間はすでに問題の条件に達しており、右の区間をくっつけたとしても条件を満たすので右の区間がある場合に+1するようにしている。
numは例題1の場合[6,1,2,7]という配列が与えられている。[6,1,2,7]からleftを動かすと[1,2,7]になりa[left]の部分はいらなくなるのでその部分をnumから引いている。

コード全体

N,K=map(int,input().split())
a=list(map(int,input().split()))
right=0
ans=0
num=0
for left in range(N):
    while num<K and right<N:
        num+=a[right]
        right+=1

    if num>=K:
        ans+=N-right+1
        num-=a[left]
print(ans)

ちなみにK以下の個数を求める場合は以下のコードになります

N,K=map(int,input().split())
a=list(map(int,input().split()))
right=0
ans=0
num=0
for left in range(N):
    while right<N and num+a[right]<=K:
        num+=a[right]
        right+=1
    ans+=right-left
    if left==right:
        right+=1
    else:
        num-=a[left]
print(ans)

少しだけ解説をいれておきます。
参考記事の問題をぱくらせていただきます

6 12
5 3 8 6 1 4

これが与えられるとします。
最初のループでleft=0,right=2までいきます。その時左の区間が[5,3]になります。この区間の合計値はK以下なので、個々の値も条件を満たすことになります。[5],[3]も条件を満たす。
上記のことを踏まえると、ansにはその区間のlenghtを足すことになるのでright-leftとしている。

参考記事
https://qiita.com/drken/items/ecd1a472d3a0e7db8dce#%E5%95%8F%E9%A1%8C-2poj-3061-subsequence