update sampling solution (#2693)

ade9bd51 · Guolin Ke · GitHub · b161f334 · ade9bd51
Unverified Commit ade9bd51 authored Jan 15, 2020 by Guolin Ke Committed by GitHub Jan 15, 2020
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

include/LightGBM/utils/random.h include/LightGBM/utils/random.h +4 -4

No files found.
--- a/include/LightGBM/utils/random.h
+++ b/include/LightGBM/utils/random.h
@@ -84,10 +84,10 @@ class Random {
      }
    } else {
      std::set<int> sample_set;
-      while (static_cast<int>(sample_set.size()) < K) {
-        int next = RandInt32() % N;
-        if (sample_set.count(next) == 0) {
-          sample_set.insert(next);
+      for (int r = N - K; r < N; ++r) {
+        int v = NextInt(0, r);
+        if (!sample_set.insert(v).second) {
+          sample_set.insert(r);
        }
      }
      for (auto iter = sample_set.begin(); iter != sample_set.end(); ++iter) {