C#

IEnumerableからGC Allocを抑えてランダム抽選する(Reservoir sampling)

この記事は約12分で読めます。

Unityでは、ちょっとした処理でも GC Alloc が発生して、GCスパイクによるフレーム落ちの原因になりがちです。
条件で絞り込んだコレクションからランダムに1つ選びたい場面で、なるべくアロケーションを発生させない方法を考えてみました。

スポンサーリンク

背景

前職ではNintendo Switchのゲーム開発をしていたので、その時考えたことをふと掘り起こして書いたので、珍しくUnityに関した記事です。
なお、この記事は Unity 標準の C# 9.0 を前提にしています。

C# compiler and language version reference • Unity Editor • Unity Docs
Details and limitations of the C# compiler and C# language version Unity uses.

長さの決まっている配列からランダムに値を取得するだけなら、乱数の最大値が明確なので、特に何も考えずに実装できます。

Hoge selected = hogeArray[rng.NextInt(hogeArray.Length)];

しかし、条件で当選対象を絞り込む場合はそうもいきません。
例えばLINQで条件付けした IEnumerable<T> からランダムに1つ取得したい場合、安直にやろうとすると、一旦 Count() するか ToArray() するかになると思います。
ただ、どちらも要素数が確定しているコレクションと比べて、余計なコストがかかってしまいます。

Countで総数を求める方法

IEnumerable<Hoge> enabled = hogeCollection.Where(hoge => hoge.IsEnabled);
int index = rng.NextInt(enabled.Count()); // ※空の場合のチェックは省略
Hoge selected = enabled.ElementAt(index);

Count() で全件を走査したあと、ElementAt() で再度途中まで走査するので、遅延評価による絞り込みが2回発生してしまいます。

ToArrayで一旦配列にする方法

Hoge[] enabledArray = hogeCollection.Where(hoge => hoge.IsEnabled).ToArray();
Hoge selected = enabledArray[rng.NextInt(enabledArray.Length)]; // ※空の場合のチェックは省略

絞り込みは1回で済みますが、要素数に比例したメモリアロケーションが発生するので、GCスパイクの原因になり得ます。

Reservoir sampling(貯水池サンプリング法)

そこで使えるのが Reservoir sampling です。
総数が分からなくても、各要素に対して1回ずつ乱数判定するだけで、一様にランダムな抽選ができます。

Reservoir sampling - Wikipedia

やることはシンプルで、i 番目の要素を確率 1/i で採用して、手元の当選候補と入れ替えていくだけです。

flowchart TD
    A["開始<br/>count = 0"] --> B{"次の要素がある?"}
    B -->|ある| C["count++"]
    C --> D{"rng.NextInt(count) == 0 ?<br/>(確率 1/count)"}
    D -->|はい| E["selected = item<br/>(当選候補を入れ替え)"]
    D -->|いいえ| F["selected はそのまま"]
    E --> B
    F --> B
    B -->|ない| G["count > 0 なら当選あり"]

乱数には Unity.Mathematics.Random を使う

乱数生成には UnityEngine.Random ではなく、Unity.Mathematics.Random を使います。

  • C#側で完結する軽量な struct なので、生成器自体もヒープ確保が無い
  • Burstとの相性も良い
  • NextInt(max) は [0, max)、NextDouble() は [0, 1) を返す
UnityEngine.Random.value は 1.0 を含む [0, 1] の範囲を返します。後述の重み付き抽選で < 1.0 の判定をすると、ちょうど 1.0 が出たときに当選漏れが起きる可能性があるので、その意味でも Unity.Mathematics.Random のほうが扱いやすいです。

実装

C#の IEnumerable<T> の拡張メソッドとしての実装例です。
Random は struct なので、ref で渡すのがポイントです(理由は後述)。

using System;
using System.Collections.Generic;
using Random = Unity.Mathematics.Random;

public static class RandomSampleExtensions
{
    public static bool TryRandomSample<T>(this IEnumerable<T> source, ref Random rng, out T selected)
    {
        if (source is null) throw new ArgumentNullException(nameof(source));

        selected = default;
        int count = 0;

        foreach (T item in source)
        {
            count++;
            // count 番目の要素を確率 1/count で採用
            if (rng.NextInt(count) == 0)
            {
                selected = item;
            }
        }

        return count > 0;
    }
}

使う側はこんな感じです。

using Random = Unity.Mathematics.Random;

public class Hoge : MonoBehaviour
{
    // readonly を付けると ref で渡せないので注意
    private Random _rng;

    private void Awake()
    {
        // シードに 0 は使えないので、下位ビットを立てておく
        _rng = new Random((uint)System.Environment.TickCount | 1u);
    }

    private void Lottery()
    {
        if (_hogeList.Where(hoge => hoge.IsEnabled).TryRandomSample(ref _rng, out Hoge selected))
        {
            // 当選処理
        }
    }
}
Unity.Mathematics.Random は struct なので、ref を付けずに値渡しすると、メソッド内のコピーの状態しか進みません。
呼び出し元の _rng は進まないままなので、何回抽選しても同じ結果が出る…という地味にハマるバグになります。
また、シードに 0 を渡すと例外になるので、こちらも注意です。

重み付けのサンプル

よくあるガチャとかで使えるやつです。
考え方は同じで、確率 w / totalWeight(その要素の重み ÷ ここまでの重みの合計)で入れ替えていきます。

public static bool TryWeightedRandomSample<T>(
    this IEnumerable<T> source,
    Func<T, double> weightSelector,
    ref Random rng,
    out T selected)
{
    if (source is null) throw new ArgumentNullException(nameof(source));
    if (weightSelector is null) throw new ArgumentNullException(nameof(weightSelector));

    selected = default;
    double totalWeight = 0;

    foreach (T item in source)
    {
        double w = weightSelector(item);
        if (!(w > 0)) continue; // 0以下と NaN の重みは無視

        totalWeight += w;
        // 確率 w / totalWeight で現在の選択を更新
        if (rng.NextDouble() < w / totalWeight)
        {
            selected = item;
        }
    }

    return totalWeight > 0;
}

最初に有効な重みが来たときは w / totalWeight がちょうど 1.0 になり、NextDouble() は1.0未満を返すので、必ず採用されます。

また、複数個を抽出したい場合も、結果を入れるバッファ(長さ k)を用意すれば同じ考え方で実現できます。
バッファを呼び出し側から渡す形にすれば、アロケーションを増やさずに済みます。

本当にゼロアロケーション?

ここまでで、要素数に比例するアロケーションは無くなりました。
とはいえ、LINQを経由している以上、完全に 0 B ではありません。

  • Where() を呼ぶたびに、イテレータオブジェクトが1つヒープに確保される
  • ラムダが外側の変数をキャプチャしていると、クロージャとデリゲートも確保される
  • List<T> を IEnumerable<T> として foreach すると、struct の Enumerator がボックス化される

いずれも要素数に関係ない小さな固定量の確保ではありますが、毎フレーム呼ぶような処理だと気になるところです。

完全ゼロアロケーション版

ProfilerのGC Allocを 0 B にしたい場合は、LINQを使わずに IReadOnlyList<T> と条件の Predicate<T> を受け取って、for で回す形にします。
List<T> も配列も IReadOnlyList<T> を実装しているので、どちらにも使えます。

public static bool TryRandomSample<T>(
    this IReadOnlyList<T> source,
    Predicate<T> match,
    ref Random rng,
    out T selected)
{
    if (source is null) throw new ArgumentNullException(nameof(source));
    if (match is null) throw new ArgumentNullException(nameof(match));

    selected = default;
    int count = 0;

    for (int i = 0; i < source.Count; i++)
    {
        T item = source[i];
        if (!match(item)) continue;

        count++;
        if (rng.NextInt(count) == 0)
        {
            selected = item;
        }
    }

    return count > 0;
}
if (_hogeList.TryRandomSample(hoge => hoge.IsEnabled, ref _rng, out Hoge selected))
{
    // 当選処理
}
ラムダ式 hoge => hoge.IsEnabled のように外部変数をキャプチャしない記述であれば、C# 9.0 (Unity) でもコンパイラがデリゲートを自動で static キャッシュしてくれるため、毎フレーム呼んでも GC Alloc は発生しません。
逆に、ローカル変数などをキャプチャすると呼び出しのたびに確保が走るので、そこだけ注意が必要です。

確率が一様であることの説明

k=1(1個だけ選ぶ)の場合で確認してみます。

ステップ $i$ で新しい要素を確率 $1/i$ で当選させ、当選したら手元の1個と入れ替えます。
要素 $j$($1 \le j \le n$)が最後まで残るには、

  • ステップ $j$ で当選する(確率 $1/j$)
  • その後のステップ $t = j+1, \dots, n$ で、入れ替えが起きない(確率 $1 – 1/t = (t-1)/t$)

必要があります。手元に残っているのは自分1個だけなので、入れ替えが起きたら必ず追い出されるんですね。

よって、

$$
\Pr[\text{要素 } j \text{ が最後に残る}]
= \frac{1}{j} \cdot \prod_{t=j+1}^{n} \frac{t-1}{t}
= \frac{1}{j} \cdot \frac{j}{n}
= \frac{1}{n}
$$

となり、どの要素も等しく $1/n$ で当選することがわかります。
(積の部分は $\frac{j}{j+1} \cdot \frac{j+1}{j+2} \cdots \frac{n-1}{n}$ と分子・分母が打ち消し合って $j/n$ になります)

重み付き版も同様で、要素 $i$ の重みを $w_i$、ステップ $t$ までの重みの合計を $W_t$ とすると、$\frac{w_i}{W_i} \cdot \prod_{t=i+1}^{n} \frac{W_{t-1}}{W_t} = \frac{w_i}{W_n}$ となり、重みに比例した確率で当選します。

図解

$n = 4$ の場合に、各ステップ終了時点で「手元に要素 $j$ が残っている確率」を表にするとこうなります。

要素1 要素2 要素3 要素4
Step 1 1 – – –
Step 2 1 × 1/2 = 1/2 1/2 – –
Step 3 1/2 × 2/3 = 1/3 1/2 × 2/3 = 1/3 1/3 –
Step 4 1/3 × 3/4 = 1/4 1/3 × 3/4 = 1/4 1/3 × 3/4 = 1/4 1/4
  • 新しく来た要素は、そのステップで確率 $1/i$ で当選する
  • すでに手元にいる要素は、入れ替えが起きない確率 $(i-1)/i$ を掛けられる

結果として、どのステップでも全要素の確率が $1/i$ に揃い、最終的に全員 $1/4$ になります👍

まとめ

  • 条件付きの IEnumerable<T> からランダム抽選するなら、Reservoir sampling で Count() や ToArray() を使わずに済む
  • 乱数は Unity.Mathematics.Random を使い、ref で渡す
  • LINQ経由だと小さな固定量の確保は残るので、完全に 0 B にしたいなら IReadOnlyList<T> + Predicate<T> 版を使う

UnityでのC#スクリプトは気を付ける点が多すぎて、ほんとにめんどくさいですね😣

以上!

コメント

タイトルとURLをコピーしました