fix Reservoir Sampling in Sample of random.h (#4450)

a06899ab · shiyu1994 · GitHub · e36cc9c1 · a06899ab
Unverified Commit a06899ab authored Jul 07, 2021 by shiyu1994 Committed by GitHub Jul 07, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 4 deletions

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

No files found.
--- a/include/LightGBM/utils/random.h
+++ b/include/LightGBM/utils/random.h
@@ -82,13 +82,10 @@ class Random {
          ret.push_back(i);
        }
      }
-    } else if (K == 1) {
-      int v = NextInt(0, N);
-      ret.push_back(v);
    } else {
      std::set<int> sample_set;
      for (int r = N - K; r < N; ++r) {
-        int v = NextInt(0, r);
+        int v = NextInt(0, r + 1);
        if (!sample_set.insert(v).second) {
          sample_set.insert(r);
        }