優先度付きキュー(priority queue)は、要素を優先度付きで追加し、要素を取り出すときは最も高い優先度を持つものから取り出されます。ただいったん格納した要素の優先度を変更したり要素がもつ値を変更することはできません。そこで今回は格納した要素の優先度を変更できる優先度つきキューを実装することにします。
優先度を変更したい要素を指定できるように要素には一意のキーをもたせます。同じキーをもつ要素を追加する場合は要素の更新として処理をします。優先度付きキューもキューである以上、格納した要素にランダムアクセスすることはできません。そこで辞書を併用することでキーから要素を取得できるようにします。また要素を削除するときも優先度がもっとも高い要素(優先度の値が最小の要素)しか削除できないので要素に死亡フラグをもたせることで任意の要素を削除できるようにします。
Elementクラスの定義
Elementクラスを定義します。フィールド変数にキー、データとしての値、優先度、死亡フラグを定義しています。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
public class UpdateablePriorityQueue<TKey, TValue> { public class Element { public Element(TKey key, TValue value, long priority) { Key = key; ; Value = value; Priority = priority; } public TKey Key; public TValue Value; public long Priority; public bool IsDead = false; } } |
UpdateablePriorityQueueクラス
UpdateablePriorityQueueクラスの本体部分です。フィールド変数とコンストラクタを示します。
最初に述べたとおり PriorityQueueとDictionaryをフィールド変数に入れています。またPriorityQueue.Dequeue()は設定したPriorityが最小の要素を取り除き、戻り値としてこれを返します。ただPriorityが最大のものを取り出したい場合もあるのでコンストラクタの引数で変更できるようにしています。
1 2 3 4 5 6 7 8 9 10 11 |
public class UpdateablePriorityQueue<TKey, TValue> { public UpdateablePriorityQueue(bool isDescendingOrder = false) { IsDescendingOrder = isDescendingOrder; } bool IsDescendingOrder = false; // 優先度の評価を逆にする Dictionary<TKey, Element> Dic = new Dictionary<TKey, Element>(); PriorityQueue<Element, long> PQ = new PriorityQueue<Element, long>(); } |
引数で指定したキーをもつ要素があるのか、その要素の値、優先度を取得するメソッドを示します。
1 2 3 4 5 6 |
public class UpdateablePriorityQueue<TKey, TValue> { public bool ContainsKey(TKey key) => Dic.ContainsKey(key); public TValue? GetValue(TKey key) => Dic.ContainsKey(key) ? Dic[key].Value : default; public long GetPriority(TKey key) => Dic.ContainsKey(key) ? Dic[key].Priority : default; } |
要素を取り除く処理を示します。ここでは要素に死亡フラグをセットして
1 2 3 4 5 6 7 8 9 10 11 |
public class UpdateablePriorityQueue<TKey, TValue> { public void Remove(TKey key) { if (Dic.ContainsKey(key)) { Dic[key].IsDead = true; Dic.Remove(key); } } } |
要素を追加する処理を示します。同じキーをもつ要素があるかもしれないのでRemoveを呼び出しています。そのあとElementオブジェクトを生成して辞書とPriorityQueueに追加しています。要素を更新するUpdateメソッドもEnqueueメソッドと同じ処理をしています。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 |
public class UpdateablePriorityQueue<TKey, TValue> { public void Enqueue(TKey key, TValue value, long priority) { if (key == null) throw new Exception("key == null は不可"); Remove(key); Element element = new Element(key, value, priority); Dic.Add(key, element); if(!IsDescendingOrder) PQ.Enqueue(element, priority); else PQ.Enqueue(element, -priority); } public void Update(TKey key, TValue value, long priority) => Enqueue(key, value, priority); } |
優先度がもっとも高い要素を取得する処理を示します。
PriorityQueue.Peek()で取得した要素にすでに死亡フラグがセットされている場合、これは削除された要素なので取り出してそのまま捨てています。死亡フラグがセットされていない要素が見つかったらそれを返します。
1 2 3 4 5 6 7 8 9 10 11 12 13 |
public class UpdateablePriorityQueue<TKey, TValue> { public Element? Peek() { while (PQ.Count > 0 && PQ.Peek().IsDead) PQ.Dequeue(); if (PQ.Count > 0) return PQ.Peek(); else return null; } } |
優先度がもっとも高い要素を取り除く処理を示します。
Peek()を実行してPriorityQueueのなかに要素がある場合はPeek()が返した要素を取り除き、それを戻り値として返します。Peek()がnullを返したりPriorityQueueが空の場合はnullを返します。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
public class UpdateablePriorityQueue<TKey, TValue> { public Element? Dequeue() { Element? element = Peek(); if (PQ.Count > 0 && element != null) { Dic.Remove(element.Key); PQ.Dequeue(); return element; } else return null; } } |
格納されている要素の数を返す処理を示します。
1 2 3 4 5 6 7 |
public class UpdateablePriorityQueue<TKey, TValue> { public int Count { get => Dic.Count; } } |
使ってみる
解法ですが、配列 の要素を大きい順に並べたときに K 番目までになるもの(Larges)とそうでないもの(Smalls)にわけます。値が更新されたらLargesの最小値とSmallsの最大値を比較します。もしSmallsの最大値のほうがLargesの最小値よりも大きい場合は所属するグループを入れ替えます。
配列の値の更新と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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
class Program { static void Main() { (int a, int b) ReadInt2() { int[] vs = ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); return (a: vs[0], b: vs[1]); } (int a, int b, int c) ReadInt3() { int[] vs = ReadLine().Split().Select(_ => int.Parse(_)).ToArray(); return (a: vs[0], b: vs[1], c: vs[2]); } (int N, int K, int Q) = ReadInt3(); // 配列の値と優先度は同じ値を用いる // Smallsは値が一番大きいものを最初に取り出せるようにしたいのでコンストラクタの引数をtrueにする UpdateablePriorityQueue<int, int> Larges = new UpdateablePriorityQueue<int, int>(); UpdateablePriorityQueue<int, int> Smalls = new UpdateablePriorityQueue<int, int>(true); for (int i = 0; i < N; i++) { if (i < K) Larges.Enqueue(i, 0, 0); else Smalls.Enqueue(i, 0, 0); } long ans = 0; for (int i = 0; i < Q; i++) { (int idx, int x) = ReadInt2(); idx--; // 配列内の値はLargesとSmallsのどちらに格納されているかを調べて x で更新する if (Larges.ContainsKey(idx)) { ans -= Larges.GetValue(idx); Larges.Update(idx, x, x); ans += x; } else if (Smalls.ContainsKey(idx)) { Smalls.Update(idx, x, x); } // Smallsの最大値のほうがLargesの最小値よりも大きい場合は入れ替え if (Smalls.Count > 0) { var min = Larges.Peek(); var max = Smalls.Peek(); if (min.Value < max.Value) { var toSmall = Larges.Dequeue(); var toBig = Smalls.Dequeue(); ans -= toSmall.Value; ans += toBig.Value; Larges.Enqueue(toBig.Key, toBig.Value, toBig.Priority); Smalls.Enqueue(toSmall.Key, toSmall.Value, toSmall.Priority); } } Console.WriteLine(ans); } } } |