累積和と尺取り法の複合問題をPythonで解いてみる-ABC172

この記事では累積和と尺取り法を駆使することで解ける問題をPythonで実装していきます。

それぞれのアルゴリズムについての説明はこの記事では省略しますが、別記事で解説していますので累積和と尺取り法についてまず知りたいという方は以下にリンクを載せておきますのでそちらをご覧ください。

Pythonで累積和を実装してみる-ABC122 - nashidos’s diary
累積和‥累積の和をとる前処理をすることによってクエリの速度を上げるアルゴリズム

Pythonでしゃくとり法（尺取り法）を実装してみる-ABC032 - nashidos’s diary
尺取り法‥左端と右端のインデックスを条件に合わせて適切に動かすことによって最適化するアルゴリズム

問題
- 問題文
- 制約
- 実装
おわりに

問題

本記事では例題として以下の問題を取り扱います。
atcoder.jp

問題文

二台の机 A, B があります。机 A には N 冊の本が、机 B には M冊の本が、それぞれ縦に積まれています。
机 A に現在上から i番目に積まれている本 $(1≤i≤N)$ は読むのに $A_i$ 分を要し、机 B に現在上から i 番目に積まれている本 $(1≤i≤M)$ は読むのに $B_i$ 分を要します。
次の行為を考えます。

本が残っている机を選び、その机の最も上に積まれた本を読んで机から取り除く。

合計所要時間が K分を超えないようにこの行為を繰り返すとき、最大で何冊の本を読むことができるでしょうか。本を読むこと以外に要する時間は無視します。

制約

1≤N,M≤200000
1≤K≤109
1≤ $A_i$ , $B_i$ ≤109
入力中の値はすべて整数

実装

問題文を読んでまず思いつくのが貪欲法ではないでしょうか。

単純に机A、Bのうちから読む時間が少なくて済む本を選んでいく方法です。

ですがもちろんこの方法では解けません。

例えばAとBが以下のような時には貪欲法では成立しません。

N = 3, M = 4, K = 200
A = [100, 1, 1]
B = [50, 50, 50, 50]

ここで活躍するのが累積和です。

先ほどの場合で累積和をとると以下のようになります。

a = [100, 101, 102]
b = [50, 100, 150, 200]

あとは読む本の数が最大になるように全ての通りを試せば出せますが今回の制約を見てみると 1≤N,M≤200000 なので間に合いません。

そこで尺取り法の出番です。

どちらか片方を固定（今回はaを固定）し左側から探索し、固定されてない方（今回で言うとb）は右側から順に探索していきます。

Aから読む本を増やしてもBから読む本の数が増えることがないので徐々に探索範囲を減らすことができるという仕組みです。

このような順序で探索を行うことで実行時間を短くすることができます。

以下が実装例です。

N, M, K = map(int,input().split())
A = list(map(int,input().split()))
B = list(map(int,input().split()))

a = [0]
b = [0]

for i in range(N):
    a.append(a[i] + A[i])

for i in range(M):
    b.append(b[i] + B[i])
    
ans, j = 0, M
for i in range(N + 1):
    if a[i] > K:
        break
    while a[i] + b[j] > K:
        j -= 1
    ans = max(ans, i + j)
print(ans)