fix a bug in ArgMaxAtK (#1196)

* fix a bug in `ArgMaxAtK`
microsoft · Jan 12, 2018 · 3a665f6 · 3a665f6
1 parent 9e223c3
commit 3a665f6
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/include/LightGBM/utils/array_args.h b/include/LightGBM/utils/array_args.h
@@ -133,7 +133,8 @@ class ArrayArgs {
     int l = start;
     int r = end - 1;
     Partition(arr, start, end, &l, &r);
-    if ((k > l && k < r) || l == 0 || r == end - 1) {
+    // if find or all elements are the same.
+    if ((k > l && k < r) || (l == start - 1 && r == end - 1)) {
       return k;
     } else if (k <= l) {
       return ArgMaxAtK(arr, start, l + 1, k);