Problem

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 at least one index j such that |i - j| <= k and nums[j] == key.

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

Examples

Example 1:

Input: nums = [3,4,9,1,3,9,5], key = 9, k = 1
Output: [1,2,3,4,5,6]
Explanation: Here, nums[2] == key and nums[5] == key.
- For index 0, |0 - 2| > k and |0 - 5| > k, so there is no j where |0 - j| <= k and nums[j] == key. Thus, 0 is not a k-distant index.
- For index 1, |1 - 2| <= k and nums[2] == key, so 1 is a k-distant index.
- For index 2, |2 - 2| <= k and nums[2] == key, so 2 is a k-distant index.
- For index 3, |3 - 2| <= k and nums[2] == key, so 3 is a k-distant index.
- For index 4, |4 - 5| <= k and nums[5] == key, so 4 is a k-distant index.
- For index 5, |5 - 5| <= k and nums[5] == key, so 5 is a k-distant index.
- For index 6, |6 - 5| <= k and nums[5] == key, so 6 is a k-distant index.
Thus, we return [1,2,3,4,5,6] which is sorted in increasing order.

Example 2:

Input: nums = [2,2,2,2,2], key = 2, k = 2
Output: [0,1,2,3,4]
Explanation: For all indices i in nums, there exists some index j such that |i - j| <= k and nums[j] == key, so every index is a k-distant index. 
Hence, we return [0,1,2,3,4].

Solution

Method 1 - Using Set

To solve this problem:

  1. Identify all indices j where nums[j] == key.
  2. For each such index j, add all indices i in the range [j - k, j + k] to a set to avoid duplicates.
  3. Convert the set to a sorted list before returning it.

Code

Java
public class Solution {
    public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
        Set<Integer> kDistantIndices = new HashSet<>();

        // Identify all indices where nums[j] == key
        for (int j = 0; j < nums.length; j++) {
            if (nums[j] == key) {
                // For each such index j, add all indices in range [j-k, j+k] to
                // the set
                int start = Math.max(0, j - k);
                int end = Math.min(nums.length - 1, j + k);
                for (int i = start; i <= end; i++) {
                    kDistantIndices.add(i);
                }
            }
        }

        // Convert the set to a sorted list
        List<Integer> result = new ArrayList<>(kDistantIndices);
        Collections.sort(result);
        return result;
    }
}
Python
def findKDistantIndices(nums, key, k):
    k_distant_indices = set()

    # Identify all indices where nums[j] == key
    for j in range(len(nums)):
        if nums[j] == key:
            # For each such index j, add all indices in range [j-k, j+k] to the set
            start = max(0, j - k)
            end = min(len(nums) - 1, j + k)
            for i in range(start, end + 1):
                k_distant_indices.add(i)

    # Convert the set to a sorted list
    return sorted(k_distant_indices)

Complexity

  • ⏰ Time complexity: O(n log n)
    • We iterate through the entire array nums exactly once to identify the indices where nums[j] == key ⇨ O(n) time
    • For each index j where nums[j] == key, we add indices in the range [j - k, j + k] to a set.
    • In the worst case, there could be (n) such key indices.
    • For each key index j, adding indices to the set could take O(k), but this is bounded by the length of the subarray size, implying an upper limit proportional to (n) because (j - k) to (j + k) is at most 2k elements, and k can be up to (n).
    • Thus, adding elements to the set takes (O(n \times k)) but seen as (O(n)) since k is a constant compared to the linear scan.
    • Sorting a list of size n takes (O(n \log n)).
  • 🧺 Space complexity: O(n)