C++で連結リストからランダムなノードを等確率で取得する方法

問題概要

単方向連結リストが与えられたとき、その中からランダムに1つのノードを選び、その値を返す getRandom() メソッドを実装することを考えます。重要な条件は、リスト内のすべてのノードが同じ確率で選ばれるという点です。

たとえば、リストが [1, 2, 3] である場合、getRandom() は 1、2、3 のいずれかをそれぞれ 1/3 の確率で返します。

アルゴリズム（リザーバーサンプリング）

リストの長さが事前にわからなくても対応できるよう、ここでは「リザーバーサンプリング」と呼ばれる手法を使用します。getRandom() メソッド内での手順は次のとおりです。

• ret := -1、len := 1、v := 先頭ノード x で初期化します。
• v が NULL になるまで以下を繰り返します。
  • rand() を len で割った余りが 0 の場合、ret に現在のノード v の値を代入します。
  • len を 1 増やします。
  • v を次のノードへ進めます。
• ループ終了後、ret を返します。

i 番目のノードが最終的に採用される確率は 1/i であり、これを繰り返すことで、長さ n のリスト上の各ノードが等確率（1/n）で選ばれることが数学的に保証されます。計算量は O(n)、追加メモリは O(1) で済むのも大きな利点です。

実装例（C++）

以下のコードで具体的な実装を確認できます。

#include <bits/stdc++.h>
using namespace std;
class ListNode{
    public:
    int val;
    ListNode *next;
    ListNode(int data){
        val = data;
        next = NULL;
    }
};
ListNode *make_list(vector<int> v){
    ListNode *head = new ListNode(v[0]);
    for(int i = 1; i<v.size(); i++){
        ListNode *ptr = head;
        while(ptr->next != NULL){
            ptr = ptr->next;
        }
        ptr->next = new ListNode(v[i]);
    }
    return head;
}
class Solution {
    public:
    ListNode* x;
    Solution(ListNode* head) {
        srand(time(NULL));
        x = head;
    }
    int getRandom() {
        int ret = -1;
        int len = 1;
        ListNode* v = x;
        while(v){
            if(rand() % len == 0){
                ret = v->val;
            }
            len++;
            v = v->next;
        }
        return ret;
    }
};
main(){
    vector<int> v = {1,7,4,9,2,5};
    ListNode *head = make_list(v);
    Solution ob(head);
    cout << (ob.getRandom());
}

入力

リストを [1,7,4,9,2,5] で初期化
getRandom() を呼び出してランダムなノードの値を取得

出力

4
9
1

実行するたびに出力は変化しますが、どのノードもほぼ均等な頻度で選ばれることが確認できます。