くろたんく雑記帳

日常とか、わんちゃんとか、機械学習とか、競プロとか、

MENU
トップ > AtCoder > Python3で解く AtCoder Beginner Contest 172 C - Tsundoku

Python3で解く AtCoder Beginner Contest 172 C - Tsundoku

AtCoder 競プロ ABC-C 累積和 しゃくとり法 二分探索

コンテストの最中は、探索するというか貪欲的に求めると勘違いしてしまった。heapqとdequeを駆使してごちゃごちゃやったが条件を詰められず。

けんちょんさんがこの問題は、しゃくとり法や累積和 +二分探索で考えたということをTweetされていたので、２つとも考えてみることにした。

概要

問題

  • {\displaystyle A, B}の机に積んある本を読むのにかかる時間のリストがある。
  • {\displaystyle A}か{\displaystyle B}のリストの左側からしか読めない。
    • スタート時を例にすると{\displaystyle A_0}を読んだら次は{\displaystyle A_1}か{\displaystyle B_0}が読める。
    • {\displaystyle A_i}の{\displaystyle i}が若いほど上に積まれているということ
  • 読んだ結果の時間の合計が{\displaystyle K}を超えないようにする。
  • 最大で何冊読める？

解くときに考えた内容

ダメパターン 1

いわゆる全探索的に行う。{\displaystyle M, N}の制約が最大200000であり、O(MN)では間に合わない。

ダメパターン 2

読む時間が少ないものから読んでいけば、本の数は多くなるだろうと考えてGreedyに解こうとすると、複雑になりすぎてダメ。（きちんと実装しきればいけるのかもだが。。。）多分時間的にも無理でしょう（計算量把握できてない。。。）

しゃくとり法
  1. 都度どれだけ読んだかを計算するのは効率が悪いので、{\displaystyle A_i, B_i}において、0冊を含めて、{\displaystyle i+1}冊読んだら合計で何分なのかというのを計算したリストを用意する。（{\displaystyle acum, bcum}とする）
  2. {\displaystyle acum}は{\displaystyle N}までか、{\displaystyle K}を超えるまで回す。
  3. {\displaystyle bcum}はおしりから見ていって、 {\displaystyle acum + bcum}が{\displaystyle K}を超えている間は{\displaystyle bcum}をお尻から見続ける。
  4. {\displaystyle acum + bcum}が{\displaystyle K}を超えなくなった時点の{\displaystyle acum}のindexと{\displaystyle bcum}のindexを合計した値がその時点の目的である最大値となるので、maxvと比較して大きい方をmaxvとする。
  5. 引き続き{\displaystyle acum}を回し続けるが、{\displaystyle bcum}を見る時のindexは直前まで見ていたところから、探索を続ける <- ここがポイント

この探索の仕方であれば、O(M+N)ということになる。そのイメージを自分なりに図示した。 つまり、最悪ケースは矢印がそれぞれ最後まで行ったとき。 {\displaystyle bcum}がストップするところは、{\displaystyle acum_i}において、{\displaystyle acum_i + bcum}が{\displaystyle K}を超えない、一番大きいbcumのindexである。{\displaystyle acum_i + bcum}が{\displaystyle K}を超えなくなるまで、{\displaystyle j}から{\displaystyle -1}していき、そのindexを見つける実装をするということ。（と思っている。）

f:id:black_tank_top:20200628143226j:plain
今回のしゃくとり法のイメージ

単調増加であることを利用して、{\displaystyle bcum}を見る時のindexは直前まで見ていたところから、探索を続けるため、{\displaystyle O(N + M)}になり、ダメパターン1と比べて大きく計算量が減る

二分探索

考え方的に、こっちの方が簡単だということに気づいた。

  1. 都度どれだけ読んだかを計算するのは効率が悪いので、{\displaystyle A_i, B_i}において、0冊を含めて、{\displaystyle i+1}冊読んだら合計で何分なのかというのを計算したリストを用意する。（{\displaystyle acum, bcum}とする）
  2. {\displaystyle acum}は{\displaystyle N}までか、{\displaystyle K}を超える手前まで回す。
  3. {\displaystyle bcum_j \leqq K - acum}を満たす最大のindexである{\displaystyle j}を返す（だだし、0番目をリストに加えることになるのでj - 1が実際の本の冊数）
  4. 現状のmaxvと {\displaystyle i + j - 1 }比較してをmaxvとする。

{\displaystyle acum}を固定して、その時とれる最大の{\displaystyle bcum}のindexをとる。こういうときに二分探索を使うと効率がいいということが勉強できた。累積和なので単調増加であり、二分探索を使う条件にも合っている。これは勉強になった。

コード

しゃくとり法
def main():
    n, m, k = map(int, input().split())
    a = [int(x) for x in input().split()]
    b = [int(x) for x in input().split()]

    acum = [0] * (n + 1)
    bcum = [0] * (m + 1)
    for i in range(n):
        acum[i + 1] = acum[i] + a[i]
    for i in range(m):
        bcum[i + 1] = bcum[i] + b[i]
        
    j = m
    maxv = 0
    for i in range(n+1):
        if acum[i] > k:
            break
        while acum[i] + bcum[j] > k:
            j -= 1
        maxv = max(maxv, i + j)
    print(maxv)

if __name__ == '__main__':
    main()
二分探索
from bisect import bisect_right

n, m, k = map(int, input().split())
a = [int(x) for x in input().split()]
b = [int(x) for x in input().split()]
maxv = 0
acum = [0] * (n + 1)
bcum = [0] * (m + 1)
for i in range(n):
    acum[i + 1] = acum[i] + a[i]
for i in range(m):
    bcum[i + 1] = bcum[i] + b[i]
for i in range(n + 1):
    if acum[i] > k:
        break
    maxv = max(maxv, i + bisect_right(bcum, k - acum[i]) - 1)
print(maxv)