てくメモ

trivial な notes

トップ > C# > 【C#】上位 k 件を得る (2) - ヒープセレクトと部分ソート

【C#】上位 k 件を得る (2) - ヒープセレクトと部分ソート

C#

前回の続き。
【C#】上位 k 件を得る (1) - ヒープソート - てくメモ


さて、上位 k 件を得ることを考える。
（※ 以下、昇順の上位のこと。降順上位の場合は Comparer を降順用にするなどすればOK）


まず単純に、全件をソートしてから k 件までの範囲を得る方法がある。

// Span<T> span;
span.Sort();
var sorted = span[..k];

標準のソートは速いし、新たに方法を用意する手間や危険もない。


ただ、すべてソートするには母数がかなり大きいという場合や、何らかの理由でソートに伴う要素の交換回数そのものを減らしたい、というときは別の方法も選択肢になる。


ここで、ヒープ構造を利用する。
参考：std::partial_sort - cppreference.com

なお、以下では前回記事で用意したメソッドを注記なく使っている。

public static void HeapSelect<T>(scoped Span<T> span, int k) => HeapSelect(span, k, Comparer<T>.Default);
public static void HeapSelect<T>(scoped Span<T> span, int k, IComparer<T> comparer)
{
    MakeHeap(span[..k], comparer);
    for (int i = k; i < span.Length; i++)
    {
        if (comparer.Compare(span[i], span[0]) < 1)
        {
            (span[0], span[i]) = (span[i], span[0]);
            ShiftDown(span[..k], comparer);
        }
    }
}

private static void ShiftDown<T>(scoped Span<T> span, IComparer<T> comparer)
{
    int length = span.Length;
    int current = 0;
    int next = 2;
    while (next < length)
    {
        if (comparer.Compare(span[next], span[next - 1]) < 1)
            --next;
        if ((comparer.Compare(span[current], span[next]) < 1) is false)
            return;

        (span[current], span[next]) = (span[next], span[current]);
        current = next;
        next = 2 * current + 2;
    }

    --next;
    if (next < length && comparer.Compare(span[current], span[next]) < 1)
        (span[current], span[next]) = (span[next], span[current]);
}

このHeapSelectメソッドは、まずspan[..k]の範囲をヒープ化し、それからspan[k..]の範囲を走査して、上位 k 件なら先頭要素と交換しヒープに均す、というような処理を行っている。

ヒープ構造の性質として先頭要素は降順最大 = 昇順下位なので、それを交換していく操作によりspan[..k]の範囲に上位を得ることができる。

// Span<int> span; // [0..16) のシャッフルされた Span
var span2 = span.ToArray().AsSpan();
var span3 = span.ToArray().AsSpan();

SortUtil.HeapSelect(span2, 4);
SortUtil.HeapSelect(span3, 8);

Console.WriteLine(span.ToListString()); // ToListString という拡張メソッドがあるものとする
Console.WriteLine(span2.ToListString());
Console.WriteLine(span3.ToListString());

// [8, 0, 7, 13, 5, 10, 1, 2, 15, 3, 12, 14, 4, 11, 6, 9]
// [3, 2, 1, 0, 13, 10, 8, 7, 15, 5, 12, 14, 4, 11, 6, 9]
// [7, 5, 6, 2, 4, 3, 1, 0, 15, 13, 12, 14, 10, 11, 8, 9]

二段目は k = 4、三段目は k = 8として、span[..k]の範囲に上位を得ることができた。


範囲を最後にソートすることで部分ソートが成立する。

public static void PartialSort<T>(scoped Span<T> span, int k) => PartialSort(span, k, Comparer<T>.Default);
public static void PartialSort<T>(scoped Span<T> span, int k, IComparer<T> comparer)
{
    HeapSelect(span, k, comparer);
    SortHeap(span[..k], comparer);
}

var spanK3 = span.ToArray().AsSpan();
var spanK10 = span.ToArray().AsSpan();

SortUtil.PartialSort(spanK3, 3);
SortUtil.PartialSort(spanK10, 10);

Console.WriteLine(spanK3.ToListString());
Console.WriteLine(spanK10.ToListString());
// [0, 1, 2, 13, 8, 10, 7, 5, 15, 3, 12, 14, 4, 11, 6, 9]
// [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 14, 13, 12, 11, 10]


部分的に○○をすることで✕✕になるというアルゴリズム、感心しがち。（個人の感想）


ソート済み k 件を得ることついては、.NET 6.0 で標準入りしたPriorityQueue（ドキュメントによれば中身は quaternary min-heap）を使う方法もある。

PriorityQueue<int, int> q = new(array.Select(v => (v, v)));
IEnumerable<TElement> dequeueK<TElement, TPriority>(PriorityQueue<TElement, TPriority> priorityQueue, int k)
{
    while (priorityQueue.TryDequeue(out var item, out _) && 0 < k--)
        yield return item;
}

Console.WriteLine(dequeueK(q, 10).ToListString());
// [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

標準ライブラリで用意されているクラスであり、自然な選択肢。

なお条件が揃わないので、この後に行う比較ではPriorityQueueは取り扱わない。


さて、BenchmarkDotNet で全件ソートして切り出すのと部分ソートとを較べてみる。
母数 N を 200 / 5,000 / 150,000、k を 10, 150 とした。

計測用コード折りたたみ

[Params(150000, 5000, 200)]
public int N;
[Params(10, 150)]
public int k;

public int[] stored = default!;
public int[] array = default!;

[GlobalSetup]
public void Setup()
{
    stored = Enumerable.Range(0, N).ToArray();
}

[IterationSetup]
public void ItrSetup()
{
    array = (int[])stored.Clone();
}

[Benchmark]
public int[] BySortAndSlice()
{
    var span = array.AsSpan();
    span.Sort();
    return span[..k].ToArray();
}

[Benchmark]
public int[] ByPartialSort()
{
    var span = array.AsSpan();
    SortUtil.PartialSort(span, k);
    return span[..k].ToArray();
}


         Method |      N |   k |         Mean |      Error |     StdDev |       Median | Allocated |
--------------- |------- |---- |-------------:|-----------:|-----------:|-------------:|----------:|
 BySortAndSlice |    200 |  10 |     3.600 μs |   6.578 μs |  0.3606 μs |     3.500 μs |     664 B |
  ByPartialSort |    200 |  10 |     9.633 μs | 112.378 μs |  6.1598 μs |     6.800 μs |     664 B |
 BySortAndSlice |    200 | 150 |     3.767 μs |   7.373 μs |  0.4041 μs |     4.000 μs |    1224 B |
  ByPartialSort |    200 | 150 |    28.167 μs |  84.756 μs |  4.6458 μs |    26.000 μs |    1224 B |
 BySortAndSlice |   5000 |  10 |    45.800 μs |  88.496 μs |  4.8508 μs |    43.100 μs |     664 B |
  ByPartialSort |   5000 |  10 |    38.333 μs |  88.182 μs |  4.8336 μs |    35.900 μs |     664 B |
 BySortAndSlice |   5000 | 150 |    44.900 μs |  85.493 μs |  4.6861 μs |    42.500 μs |    1224 B |
  ByPartialSort |   5000 | 150 |    59.733 μs |  88.333 μs |  4.8418 μs |    57.400 μs |    1224 B |
 BySortAndSlice | 150000 |  10 | 1,665.100 μs | 684.033 μs | 37.4941 μs | 1,653.200 μs |     664 B |
  ByPartialSort | 150000 |  10 |   701.900 μs |  94.797 μs |  5.1962 μs |   698.900 μs |     664 B |
 BySortAndSlice | 150000 | 150 | 1,694.167 μs | 528.515 μs | 28.9697 μs | 1,678.400 μs |    1224 B |
  ByPartialSort | 150000 | 150 |   743.600 μs | 145.504 μs |  7.9756 μs |   741.200 μs |    1224 B |


以下として〆
・基本、全件ソートして切り出し
・ただし、k, N 次第で部分ソートの効率がかなり良くなってくる