Problem Description

You are given an integer array nums (0-indexed), an integer key, and an integer k. Your task is to find all "k-distant indices" in the array.

A k-distant index is an index i where there exists at least one index j in the array that satisfies both conditions:

1. The distance between i and j is at most k: |i - j| <= k
2. The element at index j equals the key: nums[j] == key

In simpler terms, an index i is k-distant if there's at least one occurrence of key in the array within k positions from index i (either to the left or right).

For example, if nums = [3,4,9,1,3,9,5], key = 9, and k = 1:

• Index 1 is k-distant because index 2 has value 9 and |1-2| = 1 <= 1
• Index 2 is k-distant because it contains 9 itself (|2-2| = 0 <= 1)
• Index 3 is k-distant because index 2 has value 9 and |3-2| = 1 <= 1
• Index 4 is k-distant because index 5 has value 9 and |4-5| = 1 <= 1
• Index 5 is k-distant because it contains 9 itself
• Index 6 is k-distant because index 5 has value 9 and |6-5| = 1 <= 1

The function should return a list of all k-distant indices sorted in increasing order.

Intuition

The problem asks us to identify which indices are "close enough" to any occurrence of the key value in the array. The straightforward way to think about this is to check each index one by one and determine if it qualifies as k-distant.

For any given index i, we need to answer: "Is there at least one occurrence of key within distance k from this position?" To answer this, we can scan through the entire array looking for positions where nums[j] == key, and check if any of these positions satisfy |i - j| <= k.

The brute force approach naturally emerges from this line of thinking:

1. For each possible index i in the array (from 0 to n-1)
2. Check all indices j in the array to see if nums[j] == key
3. If we find even one such j where the distance |i - j| <= k, then i is k-distant
4. Add all qualifying indices to our result

This direct translation of the problem statement into code gives us a solution that's easy to understand and implement. While checking every pair of indices (i, j) might seem inefficient with O(n²) time complexity, it guarantees we won't miss any k-distant indices.

The key insight is that we don't need to find all indices j where nums[j] == key for each i - we just need to know if at least one exists within the distance constraint. This is why we can use the any() function in Python, which stops as soon as it finds the first qualifying j, providing a slight optimization in practice.

Learn more about Two Pointers patterns.

Solution Approach

The solution implements a brute force enumeration approach to find all k-distant indices.

Algorithm Steps:

1. Initialize the result list: Create an empty list ans to store all k-distant indices that we find.

2. Get array length: Store the length of the input array in variable n for convenient access.

3. Outer loop - Check each index: Enumerate through each index i from 0 to n-1. Each index i is a candidate for being k-distant.

4. Inner check - Verify k-distance condition: For each index i, we need to check if there exists at least one index j in the entire array where:

   • nums[j] == key (the element at position j equals our target key)
   • |i - j| <= k (the distance between i and j is at most k)

   This is accomplished using Python's any() function with a generator expression that iterates through all possible indices j from 0 to n-1.

5. Add qualifying indices: If the condition in step 4 is satisfied (meaning i is indeed k-distant), append i to the answer list.

6. Return sorted result: Since we iterate through indices in increasing order (0 to n-1), the resulting list is already sorted in increasing order, so we can directly return ans.

Implementation Details:

The expression any(abs(i - j) <= k and nums[j] == key for j in range(n)) efficiently checks if index i is k-distant:

• It iterates through all possible values of j
• For each j, it checks both conditions: distance constraint (abs(i - j) <= k) and value match (nums[j] == key)
• Returns True as soon as it finds the first j that satisfies both conditions
• Returns False if no such j exists after checking all indices

Time Complexity: O(n²) where n is the length of the array, as we check every pair of indices (i, j).

Space Complexity: O(1) extra space (not counting the output array), as we only use a constant amount of additional variables.

Example Walkthrough

Let's walk through the solution with a small example to illustrate how the algorithm works.

Example Input:

• nums = [2, 3, 2, 4, 3]
• key = 3
• k = 2

Step-by-step execution:

Checking index i = 0:

• We check if there's any j where nums[j] == 3 and |0 - j| <= 2
• j = 0: nums[0] = 2 ≠ 3 ❌
• j = 1: nums[1] = 3 == 3 ✓ and |0 - 1| = 1 <= 2 ✓
• Found a match! Index 0 is k-distant → Add 0 to result

Checking index i = 1:

• We check if there's any j where nums[j] == 3 and |1 - j| <= 2
• j = 0: nums[0] = 2 ≠ 3 ❌
• j = 1: nums[1] = 3 == 3 ✓ and |1 - 1| = 0 <= 2 ✓
• Found a match! Index 1 is k-distant → Add 1 to result

Checking index i = 2:

• We check if there's any j where nums[j] == 3 and |2 - j| <= 2
• j = 0: nums[0] = 2 ≠ 3 ❌
• j = 1: nums[1] = 3 == 3 ✓ and |2 - 1| = 1 <= 2 ✓
• Found a match! Index 2 is k-distant → Add 2 to result

Checking index i = 3:

• We check if there's any j where nums[j] == 3 and |3 - j| <= 2
• j = 0: nums[0] = 2 ≠ 3 ❌
• j = 1: nums[1] = 3 == 3 ✓ and |3 - 1| = 2 <= 2 ✓
• Found a match! Index 3 is k-distant → Add 3 to result

Checking index i = 4:

• We check if there's any j where nums[j] == 3 and |4 - j| <= 2
• j = 0: nums[0] = 2 ≠ 3 ❌
• j = 1: nums[1] = 3 == 3 ✓ but |4 - 1| = 3 > 2 ❌
• j = 2: nums[2] = 2 ≠ 3 ❌
• j = 3: nums[3] = 4 ≠ 3 ❌
• j = 4: nums[4] = 3 == 3 ✓ and |4 - 4| = 0 <= 2 ✓
• Found a match! Index 4 is k-distant → Add 4 to result

Final Result: [0, 1, 2, 3, 4]

All indices are k-distant because each index is within distance 2 of at least one occurrence of the key value 3 (which appears at indices 1 and 4).

Solution Implementation

Time and Space Complexity

The time complexity is O(n^2), where n is the length of the array nums. This is because we have a nested loop structure - the outer loop iterates through all n indices with for i in range(n), and for each index i, the inner comprehension any(abs(i - j) <= k and nums[j] == key for j in range(n)) iterates through all n indices again to check if there exists any index j where nums[j] == key and the distance condition abs(i - j) <= k is satisfied. This results in n * n = n^2 operations.

The space complexity is O(1) for auxiliary space (excluding the output array). The only extra variables used are loop counters and temporary values in the comprehension. However, if we consider the output array ans, the total space complexity would be O(n) in the worst case when all indices satisfy the condition and need to be included in the result.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Inefficient Time Complexity for Large Arrays

The current O(n²) solution checks every possible pair of indices, which becomes extremely slow for large arrays. Even when there are only a few occurrences of key, we still check all n indices for each position.

Example scenario: If nums has 10,000 elements but only contains the key once at index 5000, and k = 2, the current solution still performs 100,000,000 operations to find just 5 k-distant indices.

Optimized Solution:

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        # First, find all indices where key appears
        key_indices = [j for j, num in enumerate(nums) if num == key]
      
        # If no key found, return empty list
        if not key_indices:
            return []
      
        result = []
        n = len(nums)
      
        # For each index i, check if it's within k distance of any key index
        for i in range(n):
            # Use binary search or early termination for efficiency
            for j in key_indices:
                if abs(i - j) <= k:
                    result.append(i)
                    break  # Found one match, no need to check other key indices
      
        return result

2. Even Better: Range Merging Approach

A more elegant solution treats each occurrence of key as creating a "range" of k-distant indices and merges overlapping ranges:

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        key_indices = [j for j, num in enumerate(nums) if num == key]
      
        if not key_indices:
            return []
      
        result = []
        n = len(nums)
      
        # For each key position, mark the range [j-k, j+k] as k-distant
        for j in key_indices:
            start = max(0, j - k)
            end = min(n - 1, j + k)
          
            # Add all indices in this range (avoiding duplicates)
            for i in range(start, end + 1):
                if not result or result[-1] < i:
                    result.append(i)
      
        return result

3. Memory Optimization Using Set

When dealing with potential duplicates and needing sorted output, using a set can be cleaner:

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        result_set = set()
        n = len(nums)
      
        # Find all positions of key
        for j in range(n):
            if nums[j] == key:
                # Add all indices within k distance
                for i in range(max(0, j - k), min(n, j + k + 1)):
                    result_set.add(i)
      
        return sorted(result_set)

This approach is O(n*k) in the worst case but typically performs much better than O(n²) when the key appears infrequently or k is small.