ABC476を解いた【累積和・二分探索】

筆者はレート800前後の茶～緑コーダ

ABC476のD問題を解いていく

実装コード

ドリンクを買う個数を全探索し、累積和と二分探索で残りの所持金から買えるデザートの最大個数を求める。

• N, M, K と紙幣の枚数 X, Y を受け取り、価格の配列 A, B を昇順にソートする。同じ個数なら安い商品から選ぶことで、支払額と必要な紙幣の枚数を抑えられる
• 各ドリンクに必要な K ドル紙幣の枚数を (b - 1) // K + 1 で切り上げ計算し、BK に保存する
• デザートの価格の累積和 AC、ドリンクの価格の累積和 BC、ドリンクに必要な K ドル紙幣の枚数の累積和 BKC を作る
• 安い方から買うドリンクの個数 ri を 0 から M まで試し、購入後の所持金 s = X + Y * K - BC[ri] と、残る K ドル紙幣の枚数 sk = Y - BKC[ri] を求める
• s < 0 または sk < 0 なら購入できず、個数を増やしても条件を満たせないため探索を打ち切る
• ドリンクのお釣りもデザートの購入に使えるので、bisect_right(AC, s) - 1 で合計価格が s 以下になるデザートの最大個数 li を求める
• li + ri で ans を更新し、最後に購入できる商品数の最大値を出力する

main.py

from bisect import bisect_left, bisect_right, insort_left, insort_right
from collections import defaultdict, Counter, deque
from functools import reduce, lru_cache
from itertools import product, accumulate, groupby, combinations
import sys
import os
def rI(): return int(sys.stdin.readline().rstrip())
def rLI(): return list(map(int,sys.stdin.readline().rstrip().split()))
def rI1(): return (int(sys.stdin.readline().rstrip())-1)
def rLI1(): return list(map(lambda a:int(a)-1,sys.stdin.readline().rstrip().split()))
def rS(): return sys.stdin.readline().rstrip()
def rLS(): return list(sys.stdin.readline().rstrip().split())
IS_LOCAL = int(os.getenv("ATCODER", "0"))==0
err = (lambda *args, **kwargs: print(*args, **kwargs, file=sys.stderr)) if IS_LOCAL else (lambda *args, **kwargs: None)

def main():
    N, M, K = rLI()
    X, Y = rLI()
    A = sorted(rLI())
    B = sorted(rLI())

    BK = [(b - 1) // K + 1 for b in B]
    AC = [0]
    BC = [0]
    BKC = [0]
    for a in A:
        AC.append(AC[-1] + a)
    for b, bk in zip(B, BK):
        BC.append(BC[-1] + b)
        BKC.append(BKC[-1] + bk)

    ans = 0
    for ri in range(M + 1):
        s = X + Y * K - BC[ri]
        sk = Y - BKC[ri]
        if s < 0 or sk < 0:
            break
        li = bisect_right(AC, s) - 1
        ans = max(ans, li + ri)

    print(ans)

if __name__ == '__main__':
    main()

感想

所持金の合計だけでなく、ドリンクに必要な紙幣の枚数も別に管理するのがポイントだった。
ドリンクの個数を固定すると、残りは累積和と二分探索で求められるのが分かりやすい。