はじめに
前回投稿した mex のライブラリ化で、「削除機能を持った優先度付きキュー」について触れました。
前回の記事では、
- キーは、範囲が制限された非負整数
- .NET の
PriorityQueue<TElement, TPriority>クラスを利用する
という場合のソースコードを示しましたが、今回は、
- キーは、任意の
T型の値 - 優先度付きキューの主要機能である up-heap や down-heap などを自作する
という場合の実装を考えていきます。
なお、前回の記事の続きとなるため、競技プログラミング Advent Calendar 2023 の 23 日目の記事として、あとから登録しました。
実装 1
基本的な方針は前回と同様で、各要素の重複数を保持しておき、先頭の要素が削除済 (重複数が 0) になっていれば二分ヒープから削除します (遅延評価)。
任意の T 型の各要素の重複数を管理するために、連想配列 (.NET では Dictionary<T, int>) を利用します。
ソースコードは次のようになります (C#)。
// ImplicitUsings: enable
[System.Diagnostics.DebuggerDisplay("Count = {Count}")]
public class RemovableListHeapQueue<T>
{
readonly IComparer<T> c;
readonly List<T> l = new List<T>();
readonly Dictionary<T, int> counts = new Dictionary<T, int>();
int n;
public RemovableListHeapQueue(IEnumerable<T> items = null, IComparer<T> comparer = null)
{
c = comparer ?? Comparer<T>.Default;
if (items != null) foreach (var x in items) Push(x);
}
public IComparer<T> Comparer => c;
public int Count => n;
public T First => n != 0 ? l[0] : throw new InvalidOperationException("No items.");
public void Clear()
{
l.Clear();
counts.Clear();
n = 0;
}
public int GetCount(T item) => counts.GetValueOrDefault(item);
bool TrySwap(int i)
{
var p = (i - 1) / 2;
if (c.Compare(l[p], l[i]) <= 0) return false;
(l[i], l[p]) = (l[p], l[i]);
return true;
}
void UpHeap()
{
for (var i = l.Count - 1; i > 0; i = (i - 1) / 2)
{
if (!TrySwap(i)) break;
}
}
void DownHeap()
{
for (var i = 1; i < l.Count; i = i * 2 + 1)
{
if (i + 1 < l.Count && c.Compare(l[i + 1], l[i]) < 0) ++i;
if (!TrySwap(i)) break;
}
}
public void Push(T item)
{
l.Add(item);
UpHeap();
counts[item] = counts.GetValueOrDefault(item) + 1;
++n;
}
public T Pop()
{
if (n == 0) throw new InvalidOperationException("No items.");
var item = l[0];
Remove(item);
return item;
}
public bool Remove(T item)
{
if (!counts.TryGetValue(item, out var count)) return false;
if (--count == 0) counts.Remove(item); else counts[item] = count;
--n;
EnsureFirst();
return true;
}
void EnsureFirst()
{
while (l.Count > 0 && !counts.ContainsKey(l[0]))
{
l[0] = l[l.Count - 1];
l.RemoveAt(l.Count - 1);
DownHeap();
}
}
}
List<T> で二分ヒープを構成し、インデックスは 0 から使っています。
Dictionary<T, int> が重いという印象があるかもしれませんが、呼び出される回数が少ないため、性能にはほとんど影響ありません。
Pop や Remove が呼び出されたときに、最後に EnsureFirst を呼び出すことで、先頭の要素がつねに未削除の要素となるように保ちます。
時間計算量は、サイズを $n$ としたとき First は $O(1)$ 、Push は $O( \log n)$ 、そして Pop と Remove は最悪 $O(n \log n)$ ですが償却で $O( \log n)$ となります。
実装 2
ここで、EnsureFirst メソッドをもう少し軽くできないかを考えてみます。
例えば、先頭の要素のみ未削除で他の要素がすべて削除済のときに Pop を呼び出すと、計算量 $O(n \log n)$ で二分ヒープから要素が削除されていき、最終的に空になります。
そこで EnsureFirst メソッドを、末尾の要素が削除済の場合は先頭に移動させずにそのまま二分ヒープから削除するという実装に変更することで、このような例に対して $O(n)$ にできます。
void EnsureFirst()
{
while (l.Count > 0 && !counts.ContainsKey(l[0]))
{
while (l.Count > 0 && !counts.ContainsKey(l[l.Count - 1]))
l.RemoveAt(l.Count - 1);
if (l.Count == 0) break;
l[0] = l[l.Count - 1];
l.RemoveAt(l.Count - 1);
DownHeap();
}
}
PriorityQueue<TElement, TPriority> クラスを利用せずに二分ヒープを自作することで、このような処理が可能となります。
性能テスト
最大 150 万件のデータを扱う ABC 194 E - Mex Min で実行時間を比較してみます。
最も重いテストケース answer_n_00 が、実装 1 では 873 ms、実装 2 では 522 ms となりました。
また、連想配列を利用して各要素のインデックスを管理するなど、遅延評価せずに直接削除できる実装もありますが、こちらはかなり性能が落ちてしまいます。
作成したサンプル
検証したバージョン
- 開発環境: C# 10, .NET 6
- 実行環境: C# 11, .NET 7
問題集
- ABC 194 E - Mex Min
- ABC 330 E - Mex and Update
- ABC 331 E - Set Meal
- ABC 306 E - Best Performances
- ABC 281 E - Least Elements
参照
前回: mex のライブラリ化