ABC297Eを解いた【優先度付きキュー?】

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

ABC297のE問題を解いていく

実装コード

すべてのたこ焼きの値段と一番安いやつとなにか一つ買ったときの値段を優先度付きキューに入れて、値段を更新するまでキューを取り出し続ける

値段を更新したらその値段からどれか一つ買った場合すべてを優先度付きキューに入れる。

main.py

import sys
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())
def err(*args): print(*args, file=sys.stderr)

import heapq

class Minheapq(object):
    q = []
    def push(self, v):
        heapq.heappush(self.q, v)

    def pop(self):
        return heapq.heappop(self.q)

    def get(self):
        return self.q[0]

    def values(self):
        for _q in self.q:
            yield _q

def main():
    N, K = rLI()
    A = sorted(rLI())
    q = Minheapq()
    L = [float('inf')]*K
    L[0] = A[0] 
    b = A[0]
    for a in A:
        q.push(a)
        q.push(a+A[0])
    for i in range(1,K):
        while L[i-1] == (b := q.pop()):
            pass
        L[i] = b
        for a in A:
            q.push(a+b)
    print(L[-1])    
    
if __name__ == '__main__':
    main()

まとめ

公式解説は集合を使うみたいな記述をしていたが、
切ったり貼ったりを繰り返していたら
計算量が増えてむしろ遅くなったので優先度付きキューを使った。

SortedSetを使えば多分もうちょっと効率良くできる？