Back to blog
Dec 12, 2024
4 min read

Find All K-Distant Indices in an Array

Find all indices i in an array such that there exists an index j where nums[j] equals key and the absolute difference between i and j is at most k.

Difficulty: Easy | Acceptance: 77.30% | Paid: No Topics: Array, Two Pointers

You are given a 0-indexed integer array nums and two integers key and k. A k-distant index is an index i of nums for which there exists an index j such that nums[j] == key and abs(i - j) <= k.

Return the list of all k-distant indices sorted in increasing order.

Examples

Input: nums = [3,4,9,1,3,9,2], key = 9, k = 1
Output: [1,2,3,4,5,6]
Explanation: The k-distant indices with key = 9 are:
- nums[2] == 9 and abs(0 - 2) > 1
- nums[2] == 9 and abs(1 - 2) <= 1
- nums[2] == 9 and abs(2 - 2) <= 1
- nums[5] == 9 and abs(3 - 5) > 1
- nums[5] == 9 and abs(4 - 5) <= 1
- nums[5] == 9 and abs(5 - 5) <= 1
- nums[5] == 9 and abs(6 - 5) <= 1
Hence, the k-distant indices are [1,2,3,4,5,6].
Input: nums = [2,2,2,2,2], key = 2, k = 2
Output: [0,1,2,3,4]
Explanation: All indices are k-distant indices because every element is equal to key and the distance is at most 2.

Constraints

1 <= nums.length <= 1000
1 <= nums[i] <= 1000
key is an element in nums.
1 <= k <= nums.length

Brute Force

Intuition Check every index i against every index j to see if the condition nums[j] == key and abs(i - j) &lt;= k holds true.

Steps

  • Initialize an empty list result.
  • Iterate through each index i from 0 to n-1.
  • For each i, iterate through each index j from 0 to n-1.
  • If nums[j] == key and abs(i - j) &lt;= k, add i to result and break the inner loop (since we only need one valid j per i).
  • Return result.
python
class Solution:
    def findKDistantIndices(self, nums: list[int], key: int, k: int) -&gt; list[int]:
        n = len(nums)
        res = []
        for i in range(n):
            for j in range(n):
                if nums[j] == key and abs(i - j) &lt;= k:
                    res.append(i)
                    break
        return res

Complexity

  • Time: O(n²)
  • Space: O(1) (excluding output)
  • Notes: Simple to implement but inefficient for large arrays.

Boolean Array Marking

Intuition Instead of checking every i, find the valid ranges for i based on where key is located. For every index j where nums[j] == key, all indices i in the range [j-k, j+k] are valid.

Steps

  • Create a boolean array visited of size n, initialized to false.
  • Iterate through the array to find indices j where nums[j] == key.
  • For each such j, mark all indices i from max(0, j-k) to min(n-1, j+k) as true in the visited array.
  • Collect all indices i where visited[i] is true.
  • Return the collected indices.
python
class Solution:
    def findKDistantIndices(self, nums: list[int], key: int, k: int) -&gt; list[int]:
        n = len(nums)
        visited = [False] * n
        for j in range(n):
            if nums[j] == key:
                start = max(0, j - k)
                end = min(n - 1, j + k)
                for i in range(start, end + 1):
                    visited[i] = True
        return [i for i in range(n) if visited[i]]

Complexity

  • Time: O(n * k)
  • Space: O(n)
  • Notes: More efficient than brute force when k is small.

Precompute Nearest Key

Intuition For each index i, we only need to know the closest key to the left and right. We can precompute these distances in two passes.

Steps

  • Create arrays left and right of size n.
  • Iterate left to right to fill left with the index of the nearest key seen so far.
  • Iterate right to left to fill right with the index of the nearest key seen so far.
  • Iterate i from 0 to n-1. Check if abs(i - left[i]) &lt;= k or abs(i - right[i]) &lt;= k.
  • If yes, add i to result.
python
class Solution:
    def findKDistantIndices(self, nums: list[int], key: int, k: int) -&gt; list[int]:
        n = len(nums)
        left = [-1] * n
        right = [-1] * n
        
        last = -1
        for i in range(n):
            if nums[i] == key:
                last = i
            left[i] = last
            
        last = -1
        for i in range(n - 1, -1, -1):
            if nums[i] == key:
                last = i
            right[i] = last
            
        res = []
        for i in range(n):
            if (left[i] != -1 and abs(i - left[i]) &lt;= k) or (right[i] != -1 and abs(i - right[i]) &lt;= k):
                res.append(i)
        return res

Complexity

  • Time: O(n)
  • Space: O(n)
  • Notes: Most efficient time complexity, but uses extra space for the left and right arrays.