leetcode--数组中的 k 个最强值

 题目是LeetCode第192场周赛的第二题，链接：1471. 数组中的 k 个最强值。具体描述为：给你一个整数数组arr和一个整数k。设m为数组的中位数，只要满足下述两个前提之一，就可以判定arr[i]的值比arr[j]的值更强：

• |arr[i] - m| > |arr[j] - m|
• |arr[i] - m| == |arr[j] - m|，且arr[i] > arr[j]
  请返回由数组中最强的k个值组成的列表。答案可以以任意顺序返回。

 中位数是一个有序整数列表中处于中间位置的值。形式上，如果列表的长度为n，那么中位数就是该有序列表（下标从0开始）中位于((n - 1) / 2)的元素。

• 例如arr = [6, -3, 7, 2, 11]，n = 5：数组排序后得到arr = [-3, 2, 6, 7, 11]，数组的中间位置为m = ((5 - 1) / 2) = 2，中位数arr[m]的值为6。
• 例如arr = [-7, 22, 17, 3]，n = 4：数组排序后得到arr = [-7, 3, 17, 22]，数组的中间位置为m = ((4 - 1) / 2) = 1，中位数 arr[m]的值为3。

 示例1:

输入：arr = [1,2,3,4,5], k = 2
输出：[5,1]
解释：中位数为 3，按从强到弱顺序排序后，数组变为 [5,1,4,2,3]。最强的两个元素是 [5, 1]。[1, 5] 也是正确答案。
注意，尽管 |5 - 3| == |1 - 3| ，但是 5 比 1 更强，因为 5 > 1 。

 示例2:

输入：arr = [1,1,3,5,5], k = 2
输出：[5,5]
解释：中位数为 3, 按从强到弱顺序排序后，数组变为 [5,5,1,1,3]。最强的两个元素是 [5, 5]。

 示例3:

输入：arr = [6,7,11,7,6,8], k = 5
输出：[11,8,6,6,7]
解释：中位数为 7, 按从强到弱顺序排序后，数组变为 [11,8,6,6,7,7]。
[11,8,6,6,7] 的任何排列都是正确答案。

 示例4:

输入：arr = [6,-3,7,2,11], k = 3
输出：[-3,11,2]

 示例5:

输入：arr = [-7,22,17,3], k = 2
输出：[22,17]

 一看到这种求k个最大值的题目，直觉上就知道可以用一个大小为k的堆来解决，需要注意的是堆中元素大小关系的判断就需要用到题目给的条件了。一开始为了求中位数，可以用对数组排个序，然后再遍历数组，一个个都跟堆顶元素比较，更强的话入堆同时堆顶元素出堆。最后堆中的k个元素就是我们需要的了。时间复杂度为$O(nlogn+nlogk)$，空间复杂度为$O(k)$。

 JAVA版代码如下：

class Solution {
    public int[] getStrongest(int[] arr, int k) {
        int n = arr.length;
        Arrays.sort(arr);
        int m = arr[(n - 1) / 2];
        Queue queue = new PriorityQueue<>(new Comparator() {
            @Override
            public int compare(Integer i1, Integer i2) {
                if (Math.abs(i1 - m) > Math.abs(i2- m) || (Math.abs(i1 - m) == Math.abs(i2 - m) && i1 > i2)) {
                    return 1;
                }
                else {
                    return -1;
                }
            }
        });
        for (int i = 0; i < k; ++i) {
            queue.add(arr[i]);
        }
        for (int i = k; i < n; ++i) {
            int num = queue.peek();
            if (Math.abs(arr[i] - m) > Math.abs(num - m) || (Math.abs(arr[i] - m) == Math.abs(num - m) && arr[i] > num)) {
                queue.poll();
                queue.add(arr[i]);
            }
        }
        int[] result = new int[k];
        for (int i = 0; i < k; ++i) {
            result[i] = queue.poll();
        }
        return result;
    }
}

 提交结果如下:

 然后注意到其实排完序之后，最强的数肯定都排列在数组的左右两端了，使得我们可以用双指针来直接找出k个最强的值。而且因为排了序，所以左边的数肯定小于右边的数，所以题目给的判断条件里的绝对值啥的都可以去掉，甚至第二个条件就成了无用的了（可以直接去掉节省时间）。时间复杂度为$O(nlogn)$，空间复杂度为$O(1)$。

 JAVA版代码如下：

class Solution {
    public int[] getStrongest(int[] arr, int k) {
        int n = arr.length;
        Arrays.sort(arr);
        int m = arr[(n - 1) / 2];
        int left = 0, right = n - 1;
        int idx = 0;
        int[] result = new int[k];
        while (idx < k) {
            if (m - arr[left] > arr[right] - m) {
                result[idx++] = arr[left++];
            }
            else {
                result[idx++] = arr[right--];
            }
        }
        return result;
    }
}

 提交结果如下:

 Python版代码如下：

class Solution:
    def getStrongest(self, arr: List[int], k: int) -> List[int]:
        arr.sort()
        n = len(arr)
        m = arr[(n - 1) // 2]
        result = []
        count = 0
        left, right = 0, n - 1
        while count < k:
            if m - arr[left] > arr[right] - m:
                result.append(arr[left])
                left += 1
            else:
                result.append(arr[right])
                right -= 1
            count += 1
        return result

 提交结果如下: