AtCoder Beginner Contest 130 D – Enough Array をPython3で解く

Share

AtCoder上にある問題のうち、AtCoder Problemsでdiff 800以上と判定されているものを順番に解いていく企画。
基本的な考え方は全てコード中のコメントに入れてあるので、参照のこと。

出典:
AtCoder Beginner Contest 130 D – Enough Array

個人的には尺取法で書くのがどうも苦手なんだけど、頑張って書いてみた。もっとエレガントに書けそうなもんだけど……とほほ。制約が厳しくないので、二分探索で解くことも可能。

# AtCoder Beginner Contest 130 D - Enough Array
# https://atcoder.jp/contests/abc130/tasks/abc130_d
# tag: 数列 連続部分列 数え上げ 累積和 尺取法 二分探索

# 尺取法で解くのが最速。(O(N))
# [左, 右] の区間を左を中心に考え、[0, 0] から始める。
# 以下の内容をループ。

# 現在の区間が条件を満たすまで、区間右を伸ばす。
# 条件を満たしたら、そこから右の組み合わせは全て条件を満たすので、
# その数を答えに加算し、区間左を一つ縮める。
# 区間右が右端に到達しても条件を満たせなくなる or
# 区間左が右端を越える時に、ループを終了する。

def main():
N, K = map(int, input().split())
A = list(map(int, input().split()))

# 区間左、区間右、現在の合計値 [left, right)
left, right, now = 0, 0, 0
result = 0
while True:
# 条件を満たすまで区間右を伸ばす
if now < K:
if right < N:
now += A[right]
right += 1
# 伸ばしきっても条件を満たせなければ、終了
else:
break
# 条件を満たしているなら
else:
# 答えを加算 (right の直前より右側は全て条件を満たす)
result += N - right + 1
# 区間左を縮める
if left < N:
now -= A[left]
left += 1
# 区間右に追いついている場合は、区間右も一つ右側にずらす
if right < left:
now += A[right]
right += 1
# すでに区間左が右端に到達していれば、終了
else:
break

print(result)

main()

# 尺取法を使わなくても、それぞれの区間左に対する区間右を
# 二分探索で求め、回答することも可能。(O(N logN))
# その場合は、合計値を素早く求めるために、あらかじめ累積和を
# 取っておく。

# 二分探索は実装してもいいが、今回のように単純なものなら
# ライブラリ bisect を使うと早い。

from bisect import bisect_left
def main2():
N, K = map(int, input().split())
A = list(map(int, input().split()))

# 累積和を取っておく
csum = [0]
for a in A:
csum.append(csum[-1] + a)

result = 0
for left in range(N):
sum_left = csum[left]
sum_right_need = K + sum_left

# bisect による二分探索で、区間右の場所を求める
right = bisect_left(csum, sum_right_need)

# csum の最初に 0 を入れたため bisect の値も 1 ずれているので注意。
result += N - right + 1

print(result)

# main2()
Share

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です