fix Sample when sampling only one element (fix #4134) (#4324)

42ddff1d · shiyu1994 · GitHub · 4167649c · 42ddff1d
Unverified Commit 42ddff1d authored Jun 18, 2021 by shiyu1994 Committed by GitHub Jun 18, 2021
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

include/LightGBM/utils/random.h include/LightGBM/utils/random.h +3 -0

No files found.
--- a/include/LightGBM/utils/random.h
+++ b/include/LightGBM/utils/random.h
@@ -82,6 +82,9 @@ 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) {