LeetCode 382. Linked List Random Node

2022.3.10 Daily Challenge
難度：Medium

問題

這題官網說明很抽象啊。

給你一個 Linked List 的 head，請你實作 getRandom 隨機選擇一個 Node 回傳他的 val。

選擇每個 Node 的機率要一樣。

解法

直覺 Linked List 的 access 要 O(n)，單筆測資最高 call 10^4 次，Linked List 最長 10^4，時間複雜度 O(10^8) 感覺不是很樂觀，所以直接在 initializer 塞到 vector 裡，再 random 一個 int 當 vector 的 index 來隨機抽，可以 AC，但時間跟空間都只贏 10% 人左右。

class Solution
{
private:
    vector<int> v;
    function<int()> getRandInt;

public:
    Solution(ListNode *head)
    {
        ListNode *now = head;
        while (now != nullptr)
        {
            v.push_back(now->val);
            now = now->next;
        }

        random_device rd;
        mt19937 gen(rd());
        auto dis = uniform_int_distribution{0, static_cast<int>(v.size() - 1)};
        getRandInt = bind(dis, gen);
    }

    int getRandom()
    {
        return v[getRandInt()];
    }
};

後來發現這題應該是典型的 Reservoir Sampling 問題：當你想從 n 個東西裡面公平地挑 k 個（ n > k ），但你不知道 n 是多少，該怎麼挑才能使挑到任一個的機率相等。

理論這邊不提，結論來說作法是：

• 從 i = 1 開始挑，「選他」的機率要是 k/i ，寫成 pseudo-code 大概長這樣：

for (auto i = 1; ? ; i++)
{
		auto item = getItem(i);
		auto r = randomDouble(0, 1);
		if (r < (k / i))
		{
				// pick this item
		}
}

套用在這題上：

• 只要選一個「東西」，所以 k = 1
• randomDouble(0, 1) < (k / i) 其實就是 randomInt(0, i - 1) < k ，也就是 randomInt(0, i - 1) == 0
• getItem(i) 其實就是 node->val

因此解法如下：

class Solution
{
private:
    ListNode *head;

public:
    Solution(ListNode *head)
    {
        this->head = head;
    }

    int getRandom()
    {
        auto result = head->val;
        auto now = head;

        for (auto i = 1; now != nullptr; i++, now = now->next)
            if (rand() % i == 0)
                result = now->val;

        return result;
    }
};

有趣的是速度還上升了，可能 vector.push_back() 的 overhead 太大了吧。

Tags

• LeetCode