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

microsoft · Jul 7, 2021 · a06899a · a06899a
1 parent e36cc9c
commit a06899a
Showing 1 changed file with 1 addition and 4 deletions.
diff --git 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);
         }