ntk log ntk

競プロや非競技プロのやったことを書いています.

ABC123 D - Cake 123

WriterがE869120さんとsquare1001さんの双子なので別解がたくさん紹介されている.公式解法3に近いが,途中でheapqに追加していくというよりは,次の候補を先に全部追加して解いた.

D - Cake 123

解法

制約からAとBの和は全列挙できるので,まずそれを行う.AとBの和の配列ABを降順ソートし,ABの先頭の要素とCの要素それぞれとの和を優先度付きキューに入れる.このとき,和とABのindexをペアにしてから入れる.

それからK回優先度付きキューから値を取り出す.i回目の操作では,i番目に大きいAとBとCの和が取り出せる.

ABのindexをインクリメントしてから,そのABとそのとき取り出しているCと和を取り,再び優先度付きキューに戻すことを繰り返す.

pythonのheapqを扱うときは,3要素以上のタプルをheapqに入れると計算時間が長くなることが多い.公式解法3よりもこちらの解法のがよさそうな気がする.またpythonのheapqは最小ヒープなので,マイナスをつけて最大要素を取り出せるようにする.

コード

def main():
    X, Y, Z, K = (int(i) for i in input().split())
    A = [int(i) for i in input().split()]
    B = [int(i) for i in input().split()]
    C = [int(i) for i in input().split()]
    AB = []
    for a in A:
        for b in B:
            AB.append(a + b)
    AB.sort(reverse=True)
    from heapq import heappush, heappop
    ABC = []
    for c in C:
        heappush(ABC, (-(AB[0] + c), 0))  # 和, ABのindex
    ans = []
    while len(ans) < K:
        (s, idx) = heappop(ABC)
        s = -s
        ans.append(s)
        s -= AB[idx]
        idx += 1
        if idx < X*Y:
            s += AB[idx]
            heappush(ABC, (-s, idx))
    ans.sort(reverse=True)
    print(*ans, sep="\n")


if __name__ == '__main__':
    main()

提出コードへのリンク


これだけたくさん別解があると強い人はどう早く解くのか気になる.50分くらいかかったので,水diffを解くスピードをもっと上げたい.強い人の提出見に行く.

見に行った.ねてるくんさんの提出が参考になって,2回に分けて全列挙するみたいな感じだった.シンプル.