IEnumerable<T> のランダムな要素を選択する

d.hatena.ne.jp

上の記事で紹介されている「要素全体をメモリに保持せずに、ランダムな要素を選択する」アルゴリズムが面白いです。

C# で書いてみました。

public static class Extensions
{
    // source が空だったり、配列やリストの場合の対応は後述
    public static T RandomChoice<T>(this IEnumerable<T> source, Random random)
    {              
        T selected = default(T);
        int n = 1;
        foreach (var e in source)
        {
            if (random.Next(n) == 0)
                selected = e;
            n++;
        }
        return selected;
    }
}

お試しコード。

class Program
{   
    static void Main()
    {
        var xs = Enumerable.Range(0, 10).ToArray();
        var random = new Random();
        var freq = new int[xs.Length];
        for (int i = 0; i < 1000000; i++)
        {
            int x = xs.RandomChoice(random);
            freq[x]++;
        }
        
        for (int i = 0; i < freq.Length; i++)
            Console.WriteLine("{0} : {1}", i, freq[i]);
    }
}

実行結果です。

0 : 100246
1 : 99682
2 : 100206
3 : 100504
4 : 99295
5 : 100318
6 : 99000
7 : 100538
8 : 100414
9 : 99797

配列やリストへの対応

IEnumerable<T> の元が配列やリストの場合は、そもそも要素をループする必要がありません。Enumerable.Last が行っているように、IList<T> を実装しているなら、ループを回さないようにします。

要素が空かどうかのチェックも入れておきます。

public static class Extensions
{
    public static T RandomChoice<T>(this IEnumerable<T> source, Random random)
    {   
        if (source == null) throw new ArgumentNullException(nameof(source));

        var list = source as IList<T>;
        if (list != null)
        {
            if (list.Count > 0) return list[random.Next(list.Count)];
        }
        else
        {
            using (var e = source.GetEnumerator())
            {
                if (e.MoveNext())
                {
                    int n = 1;
                    T selected = default(T);
                    do
                    {
                        if (random.Next(n) == 0)
                            selected = e.Current;
                        n++;
                    } while (e.MoveNext());

                    return selected;
                }
            }
        }
        throw new InvalidOperationException("source is empty"); 
    }
}