fix a bug in `ArgMaxAtK` (#1196)

* fix a bug in `ArgMaxAtK`

fix a bug in `ArgMaxAtK` (#1196)
* fix a bug in `ArgMaxAtK`
3a665f62 · Guolin Ke · GitHub · 9e223c3a · 3a665f62
Unverified Commit 3a665f62 authored Jan 12, 2018 by Guolin Ke Committed by GitHub Jan 12, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

include/LightGBM/utils/array_args.h include/LightGBM/utils/array_args.h +2 -1

No files found.
--- a/include/LightGBM/utils/array_args.h
+++ b/include/LightGBM/utils/array_args.h
@@ -133,7 +133,8 @@ public:
    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);