Boot camp for Beginners hard 010【D – Enough Array 】をpython解説

Atcoder Problems　という競プロを勉強する上でのコンテンツがあります。URL

今回は、hard100問の内、10問目を解説していきたいと思います。

問題名　【D – Enough Array】

問題

入力

N ：数列の長さ
K ：部分数列の和の基準
A[i] ：i番目の数字

考察

今回の問題を見たとき、重要になりそうなのが、「部分数列の和」です。部分数列の和と言われたら、最初に思いつくのが累積和です。

累積和とは

累積和とは、任意の部分数列の和をそれまでの累積からの差分を使って高速で解くことのできる解法です。

ex) index 1(l)から3(r)までの部分数列を知りたい
index  0,1,2 3 4
A   = [3,5,2,4,1]
#という数列に対し

#Bという累積和を作成
index  0,1,2,3  4  5
B   = [0,3,8,10,14,15]　
#0始まりで,数列の合計を足していく

#上記で求めた累積和を使って、
#B[r+1] - B[l] 
#で求まります(rに関しては+1することに注意)

なので、求めたい区間の始点(l)と終点(r)が求まればO(1)で答えを求めることが可能

POINT1
部分数列の和を扱う問題は累積和を疑う

今回、累積和を求めた後、何となる全探索すれば問題できそうですが、実はそうは行きません。
なぜなら、N < 10^5であるため、始点と終点について全探索してしまうと、TLEしていしまいます。

ではどうするか？

今回は、尺取法を使って答えを出していきます。今回の場合の尺取法は、

• ある区間の和がKよりも小さければ、rを右に動かす(範囲を広げる)
• 大きければlを右に動かす(範囲を狭める)

これにより、試行回数は、最大でO(2*N)になるので、間に合います。

POINT2
尺取法を使えば、計算量を2*N に抑えられる

コード

import sys
def LI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def II(): return int(sys.stdin.readline())
def MI(): return map(int,sys.stdin.readline().rstrip().split())
def S(): return sys.stdin.readline().rstrip()
readline = sys.stdin.readline
sys.setrecursionlimit(10 ** 6)

N,K = MI()
A = LI()
cnt = 0
ruiseki= [0]#0番目に0を入れておく
for i in range(N):
    num = ruiseki[-1] + A[i]#一つ前の累積和+今回の足した数
    ruiseki.append(num)



l = 0#始点
r = 0#終点

cnt = [0]* (N+1)
#そこが終点になったことによって、部分数列の和がKを初めて超えた回数

while r < N+1:
    if ruiseki[r] -ruiseki[l] >= K:#累積和がK以上になったら
        cnt[r] += 1
        l += 1
        #始点を右に動かし、部分数列の範囲を広げる
    else:#累積和がK未満になったら
        r += 1
        #終点を右に動かし、部分数列の範囲を広げる
    if r > N:
        break
    

for i in range(1,N+1):
    cnt[i] += cnt[i-1]
print(sum(cnt))

まとめ

今回の問題も一見簡単そうですが、さすがD問題・・・典型解法を複数組み合わせていて手強い・・・

前の問題【B – Template Matching 】
次の問題【D – Disjoint Set of Common Divisors】