找出数组中的所有 K 近邻下标

标签: 数组 双指针

难度: Easy

给你一个下标从 0 开始的整数数组 nums 和两个整数 key 和 k 。K 近邻下标 是 nums 中的一个下标 i ，并满足至少存在一个下标 j 使得 |i - j| <= k 且 nums[j] == key 。

以列表形式返回按 递增顺序 排序的所有 K 近邻下标。

示例 1：

输入：nums = [3,4,9,1,3,9,5], key = 9, k = 1
输出：[1,2,3,4,5,6]
解释：因此，nums[2] == key 且 nums[5] == key 。
- 对下标 0 ，|0 - 2| > k 且 |0 - 5| > k ，所以不存在 j 使得 |0 - j| <= k 且 nums[j] == key 。所以 0 不是一个 K 近邻下标。
- 对下标 1 ，|1 - 2| <= k 且 nums[2] == key ，所以 1 是一个 K 近邻下标。
- 对下标 2 ，|2 - 2| <= k 且 nums[2] == key ，所以 2 是一个 K 近邻下标。
- 对下标 3 ，|3 - 2| <= k 且 nums[2] == key ，所以 3 是一个 K 近邻下标。
- 对下标 4 ，|4 - 5| <= k 且 nums[5] == key ，所以 4 是一个 K 近邻下标。
- 对下标 5 ，|5 - 5| <= k 且 nums[5] == key ，所以 5 是一个 K 近邻下标。
- 对下标 6 ，|6 - 5| <= k 且 nums[5] == key ，所以 6 是一个 K 近邻下标。
因此，按递增顺序返回 [1,2,3,4,5,6] 。

示例 2：

输入：nums = [2,2,2,2,2], key = 2, k = 2
输出：[0,1,2,3,4]
解释：对 nums 的所有下标 i ，总存在某个下标 j 使得 |i - j| <= k 且 nums[j] == key ，所以每个下标都是一个 K 近邻下标。 
因此，返回 [0,1,2,3,4] 。

提示：

  • 1 <= nums.length <= 1000
  • 1 <= nums[i] <= 1000
  • key 是数组 nums 中的一个整数
  • 1 <= k <= nums.length

Submission

运行时间: 27 ms

内存: 16.3 MB

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        n = len(nums)
        diff = [0] * (n + 10)
        for i, num in enumerate(nums):
            if num != key:
                continue
            left, right = max(0, i - k), min(n - 1, i + k)
            diff[left] += 1
            diff[right + 1] -= 1
        for i in range(1, len(diff)):
            diff[i] += diff[i - 1]

        return [i for i, num in enumerate(diff) if num]

Explain

此题解采用差分数组的方法来高效地标记那些受到关键字key影响的元素范围。首先，遍历数组nums，每当找到一个值等于key的元素时，计算出这个元素对其它元素产生影响的范围，即从i-k到i+k（需要考虑边界条件，确保不超出数组界限）。在差分数组diff中，对应位置left增加1表示开始受影响，right+1减去1表示这一区域受影响结束。之后，通过一次遍历累加差分数组diff，可以得到每个位置上的受影响累积值。最后，遍历累加后的差分数组，将所有非零元素的下标（表示这些位置受到key的影响并满足条件）收集到结果中。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        n = len(nums)  # 数组的长度
        diff = [0] * (n + 10)  # 创建差分数组，长度为n+10以处理边界情况
        for i, num in enumerate(nums):  # 遍历nums数组
            if num != key:
                continue  # 如果当前元素不是key，跳过
            left, right = max(0, i - k), min(n - 1, i + k)  # 计算影响范围，并处理边界
            diff[left] += 1  # 在左边界增加1
            diff[right + 1] -= 1  # 在右边界的下一个位置减1
        for i in range(1, len(diff)):  # 累加差分数组得到原数组的影响强度
            diff[i] += diff[i - 1]
        return [i for i, num in enumerate(diff) if num]  # 收集所有受影响的下标

Explore

在本题解中，选择在数组长度n的基础上加10主要是为了防止在进行差分数组操作时出现数组越界的情况。当元素的下标接近数组末尾时，右边界right可能会超出数组的实际长度。通过添加额外的空间（在这里是10个单位长度），可以确保即使在最极端的情况下（即k的值很大或数组末尾的元素是关键字），进行差分操作时也不会超出数组界限。这个数值10是一个保守的选择，足以处理大多数情况，但在某些特定情况下，根据k的具体大小和数组的长度，可能需要调整。

为了确保不会出现数组越界错误，我在计算影响区间的左边界（left）和右边界（right）时使用了Python的max和min函数。左边界left是通过`max(0, i - k)`计算的，这确保了即使i-k为负数（即i小于k），left的值也不会小于0，从而避免了负索引。右边界right是通过`min(n - 1, i + k)`计算的，这确保了即使i+k超出了数组的实际长度，right的值也不会超过数组的最大索引n-1。这样的处理方法可以有效防止数组越界。

差分数组diff用于在数组中高效地表示区间增减。当我们想在原数组的某个区间[a, b]增加一个值时，可以在差分数组在a位置增加该值，在b+1位置减去该值。之后，通过对差分数组进行累加操作，可以得到每个位置的实际增加值。具体到本题，每当数组nums中的元素等于key时，会在diff的left位置加1（表示从这个位置开始受到影响），在right+1位置减1（表示在这个位置之后的第一个位置结束影响）。通过累加diff数组，每个位置的值如果不是0，则表示该位置在某个区间内受到了key的影响。这样，通过遍历累加后的diff数组，就可以得到所有受key影响的位置。

在最后收集结果的过程中，使用diff数组而非原始的nums数组进行遍历是因为diff数组在之前的操作中已经累加并反映了每个位置是否受到了key的影响。在累加后的diff数组中，任何非零的值都表示相应的位置在原数组中受到了key的影响。因此，直接遍历diff数组并收集那些值非零的索引即可得到所有K近邻下标。这种方法更为高效，因为它避免了重新检查nums中每个元素的状态，而是直接利用了之前计算的结果。