reservoir_sampling.cpp

#include <iostream>
#include <vector>

#include <random>

std::random_device rd;
std::mt19937 rng(rd());

const int ITERATIONS = 300;
const int IN_SIZE = 1000000;
const int K = 900000;


std::vector<int> random_sample(std::vector<int> in, int k)
{

  std::vector<int> out(k, -1000000);
  //std::cout << "Len: " << out.size() << std::endl;
  //int need = m_N - m_curNumObsP;
  for (int i = 0; i < k; i++)
    {
      out[i] = in[i];
    }


  for (int i = k; i < in.size(); i++)
    {
      int j = rand() % i;
      //std::uniform_int_distribution<int> uni(0, i);
      //int j = uni(rng);
      if (j < k)
        {
          out[j] = in[i];
          //std::cout << out[j] <<std::endl;
        }
    }

  //for (int j = 0; j < out.size(); j++){
  //std::cout << j << ":\t" << out[j] <<std::endl;
  //}
  return out;
}

int main()
{
  srand(time(NULL));
  long int total = 0;
  std::vector<int> input;
  for (int i = 0; i < IN_SIZE; i++)
    {
      input.push_back(i);
    }

  for (int i = 0; i < 100000; i++){
    input[i] = 0; 
  }
  for (int i = 900000; i< IN_SIZE; i++){
    input[i] = 1000000;
  }

  for (int i = 0; i < ITERATIONS; i++)
    {

      std::vector<int> out = random_sample(input, K);
      for (int j = 0; j < out.size(); j++)
        {
          total += out[j];
          //std::cout << "hi " << out[j] << std::endl;
        }
    }

  std::cout << "Average: " << 1.0 * total / ITERATIONS / K << std::endl;
  std::cout << "Expected: " << 1.0 * float(IN_SIZE - 1) / 2.0 << std::endl;

  return 0;
}