【AtCoder】ABC309 – C問題について解説【Python】

そまちょブログのそまちょ（@somachob）です。

この記事は、AtCoder Beginner Contest 309 の C問題についての解説です。

C – Medicine
シミュレーション
計算量を減らす
さらに計算量を減らす
二分探索
参考

C – Medicine

入力例１をもとに解説します。

入力例１

入力を受け取ります。

薬については、リストで受け取っています。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

# 入力確認
print(N, K)
print(ab)

""" 出力結果
4 8
[[6, 3], [2, 5], [1, 9], [4, 2]]
"""

入力の受け取りについて、わからない場合は次の記事を参考にしてください。

【Python】input の使い方

シミュレーション

入力例１のとき、１日に飲む薬の数は次のようになります。

	１日目	２日目	３日目	４日目	５日目	６日目	７日目
薬１（６日間、３錠）	３	３	３	３	３	３
薬２（２日間、５錠）	５	５
薬３（１日間、９錠）	９
薬４（４日間、２錠）	２	２	２	２
合計	１９	１０	５	５	３	３	０

飲む薬の錠数を１日目から順番に計算していき、錠数が K 以下になった日にちを解答します。

たとえば次のように実装した場合、TLE になってしまうので、計算量を減らす必要があります。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

# 薬を飲む最大の日数
max_day = 0
for medicine in ab:
    max_day = max(max_day, medicine[0])

# １日目から最大の日数の翌日まで
for day in range(1, max_day + 2):
    # 飲む薬の錠数
    total = 0
    
    # 飲む薬の錠数を計算する
    for medicine in ab:
        # 飲み終えた薬のときはスキップ
        if day > medicine[0]:
            continue
        
        # 薬の錠数を足す
        total += medicine[1]
    
    # 飲む薬の錠数が K 錠以下のとき
    if total <= K:
        print(day)
        exit()

計算量を減らす

１日ずつ計算する必要はなく、各薬の飲む日数 a_i の翌日に飲む薬の錠数が変わっています。

薬１は、６日間なので、７日目で飲む薬の錠数が変わっています。
薬２は、２日間なので、３日目で飲む薬の錠数が変わっています。
薬３は、１日間なので、２日目で飲む薬の錠数が変わっています。
薬１は、４日間なので、５日目で飲む薬の錠数が変わっています。

	１日目	２日目	３日目	４日目	５日目	６日目	７日目
薬１（６日間、３錠）	３	３	３	３	３	３
薬２（２日間、５錠）	５	５
薬３（１日間、９錠）	９
薬４（４日間、２錠）	２	２	２	２
合計	１９	１０	５	５	３	３	０

このことから、各薬の飲む日数 a_i の翌日に飲む薬の錠数だけを計算するだけでいいことがわかります。

ただし、１日目に K 錠以下かを判定する必要があるので注意が必要です。

飲む薬の錠数を計算するときは、飲み終えるのが早い薬から順番に計算する必要があります。

薬１から計算すると、飲み終える翌日の７日目の錠数について計算するので、０錠で K 錠以下のため 7 が出力されてしまいます。

飲み終えるのが早い薬から計算するためにリスト ab をソートします。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

print(ab)	# => [[6, 3], [2, 5], [1, 9], [4, 2]]

# ソート
ab.sort()

print(ab)	# => [[1, 9], [2, 5], [4, 2], [6, 3]]

これを実装したのが次のコードです。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

# １日目が K 錠以下かを計算
# 飲む薬の錠数
total = 0

# 飲む薬の錠数を計算する
for medicine in ab:
    # 薬の錠数を足す
    total += medicine[1]

# 飲む薬の錠数が K 錠以下のとき
if total <= K:
    print(1)
    exit()


# 薬の飲む日数が小さい順に並び替え    
ab.sort()

# 各薬の飲む日数の翌日だけを計算
for medicine in ab:
    # 飲む日数の翌日
    day = medicine[0] + 1
    
    # 飲む薬の錠数
    total = 0
    
    # 飲む薬の錠数を計算する
    for medicine in ab:
        # 飲み終えた薬のときはスキップ
        if day > medicine[0]:
            continue
        
        # 薬の錠数を足す
        total += medicine[1]
    
    # 飲む薬の錠数が K 錠以下のとき
    if total <= K:
        print(day)
        exit()

これでもまだ、TLE です。

さらに計算量を減らす

薬の飲む錠数を毎回計算していましたが、実は、１日目の薬の錠数から、飲み終えた薬の錠数を引いていけば、薬の飲む錠数は計算できます。

１日目のすべての薬を飲むときの錠数は１９錠です。ここから、飲み終えた薬の錠数を引いていきます。

薬３は１日目に飲み終えるので、２日目は１９錠ー9錠＝１０錠
薬２は２日目に飲み終えるので、３日目は１０錠ー５錠＝５錠
薬４は４日目に飲み終えるので、５日目は５錠ー２錠＝３錠
薬１は６日目に飲み終えるので、７日目は３錠ー３錠＝０錠

	１日目	２日目	３日目	４日目	５日目	６日目	７日目
薬３（１日間、９錠）	９
薬２（２日間、５錠）	５	５
薬４（４日間、２錠）	２	２	２	２
薬１（６日間、３錠）	３	３	３	３	３	３
合計	１９	１０	５	５	３	３	０

これを実装したのが次のコードです。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

# １日目が K 錠以下かを計算
# 飲む薬の錠数
total = 0

# 飲む薬の錠数を計算する
for medicine in ab:
    # 薬の錠数を足す
    total += medicine[1]

# 飲む薬の錠数が K 錠以下のとき
if total <= K:
    print(1)
    exit()

    
# 薬の飲む日数が小さい順に並び替え    
ab.sort()

# 各薬の飲む日数の翌日だけを計算
for medicine in ab:
    # 飲む日数の翌日
    day = medicine[0] + 1
        
    # 飲み終えた薬の錠数を引く
    total -= medicine[1]
    
    # 飲む薬の錠数が K 錠以下のとき
    if total <= K:
        print(day)
        exit()

無事 AC できました。

二分探索

二分探索をすることで解くこともできます。

答えの日にちを二分探索で求めていきます。

具体的には次のように探索します。

K 錠以上の日を ng、K 錠以下の日を ok、K 錠以上の日と K 錠以下の日の真ん中を mid とします。
ng を０日目、ok を７日目（薬を飲む最大の日数の翌日）、mid を３日目とします。
mid の３日目が K 錠以下かを確認します。
３日目は K 錠以下のため、ok が３日目になります。
ng を０日目、ok を３日目、mid を１日目とします。
mid の１日目が K 錠以下かを確認します。
１日目は K 錠以上のため、ng が１日目になります。
これを ng と ok が隣同士になるまで繰り返します。

これで、条件を満たす日にちを求めることができます。

二分探索で実装したのが次のコードです。このコードでも AC できます。

# 入力
N, K = map(int, input().split())
ab = [list(map(int, input().split())) for _ in range(N)] 

# 飲む錠数が K 錠以下か判定する関数
def isOk(day):
    # 飲む薬の錠数
    total = 0
    
    # 飲む薬の錠数を計算する
    for medicine in ab:
        # 飲み終えた薬のときはスキップ
        if day > medicine[0]:
            continue
        
        # 薬の錠数を足す
        total += medicine[1]
        
    # 飲む薬の錠数が K 以下かを返す
    if total <= K:
        return True
    else:
        return False

      
# 二分探索の関数
def binarySearch(day):
    ng = 0
    ok = day
    
    # ok と ng が隣同士になるまで探索
    while abs(ok - ng) > 1:
        mid = (ok + ng) // 2
        if isOk(mid):
            ok = mid
        else:
            ng = mid

    # K 錠以下の日にちを返す
    return ok
    
    
# 薬を飲む最大の日数
max_day = 0
for medicine in ab:
    max_day = max(max_day, medicine[0])

# 探索の開始日
day = max_day + 1

# 二分探索の結果を表示
print(binarySearch(day))

参考

二分探索アルゴリズムを一般化〜めぐる式二分探索法のススメ〜 – Qiita

リンク