問題の趣旨は以下のとおりです。
配列 A のある要素の値を別の値に変更する。
そのあと降順にソートして先頭から K 番目までの総和を求めよ。
配列 A の長さは最大 5 × 10^5 であり、この処理も最大 5 × 10^5 回繰り返される。
実際に何度もソートしていては時間がかかりすぎます。時間短縮する方法はないのでしょうか？
優先度付きキューで解決
まず 配列 A の要素を大きい順に並べたときに K 番目までになるもの（Larges）とそうでないもの（Smalls）にわけます。値が更新されたらLargesの最小値とSmallsの最大値を比較します。もしSmallsの最大値のほうがLargesの最小値よりも大きい場合は所属するグループを入れ替えます。Larges の最小値とSmallsの最大値は優先度付きキューを使えばソートしなくても短時間で取得できます。
優先度付きキューはPriorityが最小のものしか取り出すことができません。また配列の要素の値を更新してしまうとPriorityも変更しなければなりませんが、普通の方法ではできません。そこで配列の要素の値が更新されたら優先度付きキュー内にある対応する要素は無効にして、新しい要素をそのつど追加していくことにします。
Solverクラスの定義
配列の要素と優先度付きキュー内に格納されている要素を関連付けるために以下のクラスと変数を定義します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
class Solver { class Data { public Data(int index, int value) { Index = index; Value = value; } public int Index = 0; // 配列の添字 public int Value = 0; // 値 public bool IsDead = false; // true なら無効 } PriorityQueue<Data, int> Larges = new PriorityQueue<Data, int>(); PriorityQueue<Data, int> Smalls = new PriorityQueue<Data, int>(); // 配列の添字から対応する優先度付きキューに格納されている要素を取得できるようにする Dictionary<int, Data> LargesFromIndex = new Dictionary<int, Data>(); Dictionary<int, Data> SmallsFromIndex = new Dictionary<int, Data>(); int K = 0; // 上位 K 番目までの総和を取得 long Sum = 0; } |
コンストラクタ
コンストラクタにおける処理を示します。
配列の最初の K 個は 何も考えずに Larges に格納します。それ以降は Larges の最小値と比較して条件分岐させます。追加しようとしている要素が Larges の最小値より大きい場合は Larges の最小値を Larges から削除して Smalls へ移動させ、新しく追加する要素を Larges に追加します。そうでない場合は Smalls へ追加します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
class Solver { public Solver(int[] arr, int k) { K = k; for (int i = 0; i < arr.Length; i++) { Data data = new Data(i, arr[i]); if (K > i) // 最初の K 個は Larges に格納する { Larges.Enqueue(data, data.Value); LargesFromIndex.Add(i, data); Sum += data.Value; } else // それ以降は Larges の最小値と比較して条件分岐させる { Data last = Larges.Peek(); if (last.Value < data.Value) // Larges の最小値より大きい場合は入れ替え { ToSmalls(); Larges.Enqueue(data, data.Value); LargesFromIndex.Add(i, data); Sum += data.Value; } else { Smalls.Enqueue(data, -data.Value); // Smalls は最大値を取得したいので優先度の符号を反転させる SmallsFromIndex.Add(i, data); } } } } } |
上位グループと下位グループの入れ替え
ToSmallsメソッドはLargesの最小値をSmallsに移動させ、ToLargesメソッドはSmallsの最大値をLargesに移動させます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
class Solver { void ToSmalls() { Data removed = Larges.Dequeue(); Smalls.Enqueue(removed, -removed.Value); LargesFromIndex.Remove(removed.Index); SmallsFromIndex.Add(removed.Index, removed); Sum -= removed.Value; } void ToLarges() { Data removed = Smalls.Dequeue(); Larges.Enqueue(removed, removed.Value); SmallsFromIndex.Remove(removed.Index); LargesFromIndex.Add(removed.Index, removed); Sum += removed.Value; } } |
上位グループと下位グループ内の更新
配列の要素が更新されたときは対応する要素が Larges と Smalls のどちらにあるのかを調べて見つかった要素を無効にします。そして新しい値がセットされた要素を追加しなおします。
そのあとLargesの最小値と Smallsの最大値を比較します。最大値と最小値はPriorityQueue.Peek()を実行すれば取得できるのですが、取得されたものが無効なものであれば有効なものが見つかるまで取り除きます。比較対象が見つかったら比較して必要であれば入れ替えの処理をおこないます。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
class Solver { public void Update(int index, int value) { if (LargesFromIndex.ContainsKey(index)) UpdateLarges(index, value); else UpdateSmalls(index, value); // Largesの最小値と Smallsの最大値を比較する // 最大値と最小値が無効であれば有効なものが見つかるまで取り除く while (Larges.Count > 0 && Larges.Peek().IsDead) Larges.Dequeue(); while (Smalls.Count > 0 && Smalls.Peek().IsDead) Smalls.Dequeue(); // 比較対象が存在しない if (Smalls.Count == 0) return; Data largesMin = Larges.Peek(); Data smallsMax = Smalls.Peek(); if (largesMin.Value < smallsMax.Value) { ToSmalls(); ToLarges(); } } } |
処理をしたあと上位 K 番目までの総和を取得します。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
class Solver { public long GetSum() { return Sum; } // 蛇足 上位 K 番目の値を取得する public int GetValue() { Data largesMin = Larges.Peek(); return largesMin.Value; } } |
仕上げ
あとは上記で定義したSolverクラスで解を求めるだけです。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
class Program { static void Main() { int[] ReadIntArray() => Console.ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); int[] vs1 = ReadIntArray(); int N = vs1[0]; int K = vs1[1]; int Q = vs1[2]; int[] A = new int[N]; Solver Solver = new Solver(A, K); for (int i = 0; i < Q; i++) { int[] vs2 = ReadIntArray(); int x = vs2[0] - 1; int y = vs2[1]; Solver.Update(x, y); Console.WriteLine(Solver.GetSum()); } } } |
