D - 高橋君の価格調整 / Takahashi's Price Adjustment Editorial by admin
gemini-3.5-flash-thinking概要
この問題は、\(N\) 個の商品に対して「現在最も価格が高い(または低い)商品を選択し、その価格を更新する(あるいは更新を取り消す)」というクエリを \(Q\) 回処理し、最終的な全商品の価格の合計を求める問題です。
優先度付きキュー(ヒープ)と遅延評価(Lazy Deletion)を組み合わせることで、各クエリを高速に処理します。
考察
素朴なアプローチ
毎クエリごとに配列全体を走査して最大値や最小値を探すと、1回のクエリに \(O(N)\) の時間がかかります。クエリ数は \(Q\) 回あるため、全体の計算量は \(O(NQ)\) となり、制約 \(N, Q \leq 2 \times 10^5\) では実行時間制限(TLE)になってしまいます。
したがって、最大値と最小値を高速に取得・更新できるデータ構造が必要です。
効率的なアプローチの模索
最大値・最小値を高速に管理するデータ構造として優先度付きキュー(ヒープ)が挙げられます。 最大値管理用のヒープ(最大ヒープ)と、最小値管理用のヒープ(最小ヒープ)の2つを用意すれば、それぞれ \(O(1)\) で極値を取得できます。
しかし、価格が更新されたとき、ヒープ内にある「更新前の古い価格データ」を直接削除しようとすると \(O(N)\) の時間がかかってしまいます。
遅延評価(Lazy Deletion)による解決
ヒープから古いデータを直接削除する代わりに、「取り出したデータが最新のものかどうかをその都度判定し、古ければ捨てる」という遅延評価の手法を用います。
具体的には、各商品の「現在の正しい価格」を記録する配列 P を用意しておきます。
ヒープから要素 (価格, 商品番号) を取り出した際、その 価格 が P[商品番号] と一致しているかを確認します。
- 一致している場合:最新の正しいデータなので、これを処理対象とします。
- 一致していない場合:過去に更新される前の古いデータなので、何もせず破棄し、次の要素をヒープから取り出します。
この工夫により、古いデータを無理に消しにいく必要がなくなり、高速なクエリ処理が可能になります。
アルゴリズム
初期化
- 現在の各商品の価格を保持する配列 \(P\) を \(A\) で初期化します。
- 全商品の合計金額
total_sumを \(\sum A\) で初期化します。 - 最大ヒープ
max_heapに、各商品 \(i\) について(-A[i], i)を追加します(Pythonのheapqは最小ヒープなので、値をマイナスにすることで最大ヒープとして扱います)。 - 最小ヒープ
min_heapに、各商品 \(i\) について(A[i], i)を追加します。
クエリの処理 各クエリ \((T_j, X_j, R_j)\) に対して、以下の操作を行います。
\(T_j = 1\)(最大値の選択)の場合:
max_heapから要素(neg_p, i)をポップし続けます。-neg_p == P[i](現在の価格と一致)となる要素が見つかった時点でループを抜け、その商品 \(i\) を対象とします。- \(R_j = 0\)(確定)のとき:
total_sumに \(X_j - P[i]\) を加算します。- \(P[i]\) を \(X_j\) に更新します。
- 新しい価格情報を
max_heapに(-X_j, i)、min_heapに(X_j, i)として追加します。
- \(R_j = 1\)(取り消し)のとき:
- 価格は変更されないため、取り出した最新の情報をそのまま
max_heapに戻します((-P[i], i)をプッシュ)。
- 価格は変更されないため、取り出した最新の情報をそのまま
\(T_j = 2\)(最小値の選択)の場合:
- 同様の操作を
min_heapを用いて行います。
- 同様の操作を
最終出力 すべてのクエリを処理し終えた後の
total_sumを出力します。
計算量
時間計算量: \(O((N + Q) \log N)\) ヒープに追加される要素の総数は、初期状態の \(N\) 個と、各クエリで追加される高々 \(2\) 個(\(R_j = 0\) のときに最大・最小ヒープへ1個ずつ、または \(R_j = 1\) のときに1個戻す)を合わせて、最大でも \(O(N + Q)\) 個です。 ヒープのポップ操作も最大で追加された回数分しか行われないため、1回あたりのヒープ操作 \(O(\log N)\) を考慮すると、全体の時間計算量は \(O((N + Q) \log N)\) となり、実行時間制限に余裕で間に合います。
空間計算量: \(O(N + Q)\) 配列 \(P\) に \(O(N)\)、ヒープに最大で \(O(N + Q)\) 個の要素を格納するため、空間計算量は \(O(N + Q)\) となります。
実装のポイント
タイブレーク(商品番号が最も小さいものを選ぶルール) 問題文には「価格が最も高い(低い)商品が複数あるときは、その中で商品番号が最も小さいものを選ぶ」というルールがあります。 Pythonのタプルの比較は、第1要素(価格)が同じであれば、第2要素(インデックス \(i\))の比較に移行します。
最大ヒープ:
(-価格, i)とすることで、価格が最大のものが最小の-価格となり最優先され、価格が同じなら \(i\) が小さいものが優先されます。最小ヒープ:
(価格, i)とすることで、価格が最小のものが最優先され、価格が同じなら \(i\) が小さいものが優先されます。 このようにタプルを設計することで、ヒープの標準の挙動だけで問題のタイブレーク条件を自動的に満たすことができます。ソースコード
import sys
import heapq
def solve():
input_data = sys.stdin.read().split()
if not input_data:
return
N = int(input_data[0])
Q = int(input_data[1])
A = [int(x) for x in input_data[2:2+N]]
P = list(A)
total_sum = sum(A)
max_heap = [(-val, i) for i, val in enumerate(A)]
min_heap = [(val, i) for i, val in enumerate(A)]
heapq.heapify(max_heap)
heapq.heapify(min_heap)
heappush = heapq.heappush
heappop = heapq.heappop
it = iter(input_data[2+N:])
for _ in range(Q):
t = int(next(it))
x = int(next(it))
r = int(next(it))
if t == 1:
while max_heap:
neg_p, i = heappop(max_heap)
if P[i] == -neg_p:
old_p = -neg_p
break
if r == 0:
P[i] = x
heappush(max_heap, (-x, i))
heappush(min_heap, (x, i))
total_sum += x - old_p
else:
heappush(max_heap, (-old_p, i))
else:
while min_heap:
p, i = heappop(min_heap)
if P[i] == p:
old_p = p
break
if r == 0:
P[i] = x
heappush(max_heap, (-x, i))
heappush(min_heap, (x, i))
total_sum += x - old_p
else:
heappush(min_heap, (old_p, i))
print(total_sum)
if __name__ == '__main__':
solve()
この解説は gemini-3.5-flash-thinking によって生成されました。