徒然

思いついたら書きます

可変長配列からランダムに要素を取り出したいときのTips

はじめに

Pythonならlist、c++ならvectorのような可変長配列からランダムに要素を取り出したいときはどうすればよいでしょうか?
ランダムに要素を選んで取り出したいところですが、可変長配列から要素を取り出す操作はO(n)かかってしまいます。

そんなときは、末尾の取り出しがO(1)なことを利用して、末尾と要素を入れ替えてから取り出すと高速になります。

Pythonコード例

import random
random.seed(0)
a = [random.randrange(100) for _ in range(5)]
while a:
    index = random.randrange(len(a))
    a[index], a[-1] = a[-1], a[index]
    x = a.pop()
    print(x)

出力

33
5
97
53
49

C++コード例

#include <bits/stdc++.h>
using namespace std;
int main()
{
  vector<int> a(5);
  for (int i = 0; i < int(a.size()); i++)
  {
    a[i] = rand() % 100;
  }
  while (!a.empty())
  {
    int index = rand() % a.size();
    swap(a[index], a.back());
    cout << a.back() << endl;
    a.pop_back();
  }
}

出力

69
34
41
0
67

おわりに

私がAtCoder Heuristic Contestでよく利用するテクニックでした。