Official

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[商品番号] と一致しているかを確認します。 - 一致している場合:最新の正しいデータなので、これを処理対象とします。 - 一致していない場合:過去に更新される前の古いデータなので、何もせず破棄し、次の要素をヒープから取り出します。

この工夫により、古いデータを無理に消しにいく必要がなくなり、高速なクエリ処理が可能になります。


アルゴリズム

  1. 初期化

    • 現在の各商品の価格を保持する配列 \(P\) を \(A\) で初期化します。
    • 全商品の合計金額 total_sum を \(\sum A\) で初期化します。
    • 最大ヒープ max_heap に、各商品 \(i\) について (-A[i], i) を追加します(Pythonの heapq は最小ヒープなので、値をマイナスにすることで最大ヒープとして扱います)。
    • 最小ヒープ min_heap に、各商品 \(i\) について (A[i], i) を追加します。
  2. クエリの処理 各クエリ \((T_j, X_j, R_j)\) に対して、以下の操作を行います。

    • \(T_j = 1\)(最大値の選択)の場合:

      1. max_heap から要素 (neg_p, i) をポップし続けます。
      2. -neg_p == P[i](現在の価格と一致)となる要素が見つかった時点でループを抜け、その商品 \(i\) を対象とします。
      3. \(R_j = 0\)(確定)のとき:
        • total_sum に \(X_j - P[i]\) を加算します。
        • \(P[i]\) を \(X_j\) に更新します。
        • 新しい価格情報を max_heap に (-X_j, i)、min_heap に (X_j, i) として追加します。
      4. \(R_j = 1\)(取り消し)のとき:
        • 価格は変更されないため、取り出した最新の情報をそのまま max_heap に戻します((-P[i], i) をプッシュ)。
    • \(T_j = 2\)(最小値の選択)の場合:

      • 同様の操作を min_heap を用いて行います。
  3. 最終出力 すべてのクエリを処理し終えた後の 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 によって生成されました。