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

microsoft · Jun 18, 2021 · 42ddff1 · 42ddff1
1 parent 4167649
commit 42ddff1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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) {