Merge pull request #56 from swift-nav/random_without_replacement_bug

Bug: randint_without_replacement

Author: akleeman

albatross/random_utils.h

@@ -17,14 +17,39 @@
 #include <random>
 #include <set>
 
+namespace albatross {
 /*
  * Samples integers between low and high (inclusive) with replacement.
  */
 inline std::vector<std::size_t>
 randint_without_replacement(std::size_t n, std::size_t low, std::size_t high,
                             std::default_random_engine &gen) {
   assert(n >= 0);
-  assert(n <= (high - low));
+
+  std::size_t n_choices = high - low + 1;
+  assert(n <= n_choices);
+
+  if (n == (high - low + 1)) {
+    std::vector<std::size_t> all_inds(n);
+    std::iota(all_inds.begin(), all_inds.end(), 0);
+    return all_inds;
+  }
+
+  if (n > n_choices / 2 + 1) {
+    // Since we're trying to randomly sample more than half of the
+    // points it'll be faster to randomly sample which points we
+    // should throw out than which ones we should keep.
+    const auto to_throw_out =
+        randint_without_replacement(n_choices - n, low, high, gen);
+    auto to_keep = indices_complement(to_throw_out, high - low);
+
+    if (low != 0) {
+      for (auto &el : to_keep) {
+        el += low;
+      }
+    }
+    return to_keep;
+  }
 
   std::uniform_int_distribution<std::size_t> dist(low, high);
   std::set<int> samples;
@@ -33,5 +58,6 @@ randint_without_replacement(std::size_t n, std::size_t low, std::size_t high,
   }
   return std::vector<std::size_t>(samples.begin(), samples.end());
 }
+}
 
 #endif

tests/test_random_utils.cc

@@ -33,4 +33,9 @@ TEST(test_random_utils, randint_without_replacement) {
   }
 }
 
+TEST(test_random_utils, randint_without_replacement_full_set) {
+  std::default_random_engine gen;
+  const auto inds = randint_without_replacement(10, 0, 9, gen);
+}
+
 } // namespace albatross