数组中的 k 个最强值

标签: 数组 双指针 排序

难度: Medium

给你一个整数数组 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]

提示：

  • 1 <= arr.length <= 10^5
  • -10^5 <= arr[i] <= 10^5
  • 1 <= k <= arr.length

Submission

运行时间: 129 ms

内存: 27.4 MB

class Solution:
    def getStrongest(self, arr: List[int], k: int) -> List[int]:
        l = len(arr)
        m_idx = (l - 1) // 2
        if k == l:
            return arr
        elif k > m_idx:
            reverse_flag = True # find the (l-k) weakest values
            k = l - k
        else:
            reverse_flag = False

        arr = sorted(arr)
        median = arr[m_idx]

        if reverse_flag:
            count = 0
            left_idx = m_idx - 1
            right_idx = m_idx
            left_val = median - arr[left_idx]
            right_val = arr[right_idx] - median
            while count < k:
                if left_val <= right_val:
                    left_idx -= 1
                    left_val = median - arr[left_idx]
                else:
                    right_idx += 1
                    right_val = arr[right_idx] - median
                count += 1
            return arr[:left_idx + 1] + arr[right_idx:]
        else:
            out = []
            count = 0
            left_idx = 0
            right_idx = l-1
            left_val = median - arr[left_idx]
            right_val = arr[right_idx] - median
            while count < k:
                if left_val <= right_val:
                    out.append(arr[right_idx])
                    right_idx -= 1
                    right_val = arr[right_idx] - median
                else:
                    out.append(arr[left_idx])
                    left_idx += 1
                    left_val = median - arr[left_idx]
                count += 1
            return out

Explain

此题解首先通过对数组进行排序来找到中位数m，然后根据给定的规则（与中位数的绝对差值和元素本身的大小）来确定元素的“强度”。排序后，题解使用了两个指针，从中间向两端扫描，选择最强的k个元素。如果k大于中位数索引m_idx，意味着需要找最弱的元素（实际上就是除去最强的一些元素），所以有一个reverse_flag标志来决定是寻找最强还是最弱的元素。通过比较左右指针指向的元素与中位数的绝对差值大小来决定移动哪一个指针，并记录已选元素个数，直到达到k个。

时间复杂度: O(n log n)

空间复杂度: O(n)

class Solution:
    def getStrongest(self, arr: List[int], k: int) -> List[int]:
        l = len(arr)
        m_idx = (l - 1) // 2  # 计算中位数的索引
        if k == l:
            return arr  # 如果k等于数组长度，直接返回原数组
        elif k > m_idx:
            reverse_flag = True  # 如果k大于中位数索引，寻找最弱的元素
            k = l - k  # 更新k为最弱元素的数量
        else:
            reverse_flag = False  # 否则寻找最强的元素

        arr = sorted(arr)  # 对数组排序以找到中位数
        median = arr[m_idx]  # 中位数

        if reverse_flag:
            count = 0
            left_idx = m_idx - 1
            right_idx = m_idx
            left_val = median - arr[left_idx]  # 计算左指针的强度差
            right_val = arr[right_idx] - median  # 计算右指针的强度差
            while count < k:
                if left_val <= right_val:
                    left_idx -= 1
                    left_val = median - arr[left_idx]
                else:
                    right_idx += 1
                    right_val = arr[right_idx] - median
                count += 1
            return arr[:left_idx + 1] + arr[right_idx:]  # 返回最弱的k个元素
        else:
            out = []
            count = 0
            left_idx = 0
            right_idx = l-1
            left_val = median - arr[left_idx]
            right_val = arr[right_idx] - median
            while count < k:
                if left_val <= right_val:
                    out.append(arr[right_idx])
                    right_idx -= 1
                    right_val = arr[right_idx] - median
                else:
                    out.append(arr[left_idx])
                    left_idx += 1
                    left_val = median - arr[left_idx]
                count += 1
            return out  # 返回最强的k个元素

Explore

在题解中，`reverse_flag`的使用是基于需要选择最强还是最弱的元素来确定的。如果`k`大于中位数的索引`m_idx`（即中位数左侧的元素个数），则表示选取的元素数量超过了数组中位数左侧的元素数量，因此除了选择最强的元素外，还需要从数组的右侧选择相应数量的最弱元素，以确保总共选择了`k`个元素。这时设置`reverse_flag`为真，并将`k`更新为最弱元素的数量（数组长度减去`k`），来反向选择最弱的元素。如果`k`小于或等于`m_idx`，则不需要这样的逻辑，直接寻找最强的`k`个元素即可。

在计算中位数索引时使用`(l - 1) // 2`是为了确保在数组长度为奇数或偶数时，中位数都能正确地表示为中间的元素（或左侧中间元素）。当数组长度为奇数时，`(l - 1) // 2`和`l // 2`得到同一个结果，即中间的元素。但当数组长度为偶数时，`(l - 1) // 2`得到的是左侧中间元素的索引，这与统计学中的中位数定义（在偶数个数的数据中，中位数通常取中间两个数的平均值）有所不同，但在此算法中，选择其中一个作为中位数可以简化计算。这种方式主要是为了在选择最强或最弱元素时，可以更简便地处理数组的两端。

题解中通过使用双指针策略来确保始终能选出`k`个符合条件的元素。当不使用`reverse_flag`时（即寻找最强的元素），两个指针分别从数组的两端开始，根据元素与中位数的绝对差值的大小来决定哪个指针向中间移动，这样每次可以从最强的元素中选择一个加入结果数组。当使用`reverse_flag`时（即寻找最弱的元素），逻辑类似，但是指针移动的方向相反，从中间向两边扫描，选择最弱的元素。这种方法通过逐步减少选择范围，确保每次都能找到当前最强或最弱的元素，直到收集到足够的`k`个元素。