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

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

ABC234のD問題を解いていく

実装コード

優先度付きキュー heapq を用いて、常に K 個だけ要素を保持する。

このとき、

• 優先度付きキューに含まれる要素のうち 最小値 が答えになる
• heapq では最小値が常に先頭要素として管理される

各要素を順に追加し、
要素数が K を超えた場合は、その都度最小要素を取り除く。

こうして、これまでに見た要素のうち大きい方から K 個
だけが優先度付きキューに残るため、要素数が K に達した時点で先頭を出力すればOK。

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)

import heapq

def main():
    N, K = rLI()
    P = rLI()
    A = []
    
    hq = []
    
    for i in range(N):
        heapq.heappush(A, P[i])
        if len(A) > K:
            heapq.heappop(A)
        if len(A) == K:
            print(A[0])
    
if __name__ == '__main__':
    main()

感想

二分探索用ライブラリ（insort_left）を使った案を GPT に分析させたところ、計算量が (O(N^2)) になると指摘された。

その代替として heapq を使う方法を提案され、
実際にそれを実装して公式解説を確認したところ、内容はほぼ同じだった。

この経験から、データ構造や標準ライブラリの内部仕様・計算量を正しく理解した上で選択することが不可欠だと強く実感した。