Generating multinomial random variates via the binomial conditional method

multinomial.cpp

#include <vector>
#include <numeric>
#include <random>

using namespace std;

template <typename T>
T sum(vector<T> v) {
  return accumulate(v.begin(), v.end(), 0.0);
}

template <typename Int = int>
vector<Int> rmultinom(const Int N = 100, const Int K = 4,
                      const vector<double> p = {0.25, 0.25, 0.3, 0.2}) {
  if (p.size() != K) {
    throw length_error("'p' should be of size K");
  }
  double norm = sum(p);
  if (norm != 1.0) {
    throw logic_error("'p' must sum to 1");
  }

  vector<Int> n(K);
  double sum_p = 0.0;
  Int sum_n = 0;
  mt19937 rng;
  rng.seed(random_device()());

  for (Int k = 0; k < K - 1; k++) {
    if (p[k] > 0.0) {
      // This becomes unstable when p ~= 1 - so we just take Bin(N, 1) = N
      binomial_distribution<Int> binom(N - sum_n, p[k] / (norm - sum_p));
      n[k] = binom(rng);
    } else {
      n[k] = 0;
    }
    sum_p += p[k];
    sum_n += n[k];
  }
  n[K - 1] = N - sum_n;
  return n;
}