Problem description

A naive method

1. find all the numbers that has a value of key, save their positions in a list
2. loop through the list as j
3. find index i that is not k steps further than j
4. if there is at least 1 index i that fits the requiremnt, add j into res
5. return res

Solution 1: Scan twice

Time Complexity: O(n + m × k)

Where:

• n = length of nums array
• m = number of indices where nums[i] == key
• k = the distance parameter

Breakdown:

1. First loop (building nums_pos dictionary): O(n)
   • Iterates through all n elements once
2. Nested loops (finding k-distant indices): O(m × k)
   • Outer loop: iterates through m positions where nums[i] == key
   • Inner loop: for each key position j, iterates through range [j-k, j+k], which is at most 2k+1 elements
   • Total: m × (2k+1) = O(m × k)
3. Converting set to list: O(result size) ≤ O(n) Overall: O(n + m × k)

Space Complexity: O(n)

• nums_pos dictionary: O(n) in worst case
• res set: O(n) in worst case (when all indices are k-distant)

Edge cases for complexity:

• Best case: No elements equal key → O(n)
• Worst case: All elements equal key and k is large → O(n × k)
• When k ≥ n, this effectively becomes O(n²) since m ≤ n

The algorithm is efficient for reasonable values of k, but could become expensive if k approaches n and many elements equal the key value.

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        # 1. find all the numbers that has a value of key, save their positions in a list
        # 2. loop through the list as j 
        # 3. find index i that is not k steps further than j 
        # 4. if there is at least 1 index i that fits the requiremnt, add j into res
        # 5. return res
        nums_pos = {}
        n = len(nums)
        for i,num in enumerate(nums):
            if num not in nums_pos:
                nums_pos[num] = [i]
            else:
                nums_pos[num].append(i)
        res = set()
        num_key_pos = nums_pos[key]
        for j in num_key_pos:
            left = max(0,j-k) # j-k to j, difference <= k
            right = min(j+k+1, n) # j to j+k, difference <= k
            for idx in range(left, right):
                res.add(idx)
        return list(res)

Solution 2: Scan once by tracking the next unprocessed index

Key ideas

• Tracking the next unprocessed index (r):
  • Ensures no index is added twice.
  • Automatically merges overlapping ranges.
• Guaranteeing sorted output:
  • Since indices are processed in order, the result is naturally sorted.

Time Complexity: O(n)

Key insight: The variable l ensures each index is processed at most once across all iterations.

Breakdown:

1. Outer loop: O(n) - iterates through all n elements
2. Inner loop: Although nested, it's amortized O(1) per outer iteration
   • Each index i in nums can only be added to res once
   • The variable l tracks the next unprocessed position, preventing duplicate processing
   • Total work across all inner loops = O(n) since we visit each index at most once

Why it's O(n) and not O(n²):

• Without the l = right optimization, overlapping ranges would cause duplicate processing
• With l = right, each position 0 to n-1 is processed exactly once across all iterations
• The inner loop work is distributed across the outer loop iterations

Space Complexity: O(1) auxiliary space, O(n) total space

• Auxiliary space: O(1) - only using variables l, left, right, etc.
• Output space: O(n) - res list can contain up to n elements in worst case
• Total space: O(n) (counting the output)

Example demonstrating the optimization:

nums = [3,4,9,1,3,9,5], key = 9, k = 1

• When j=2 (nums[2]=9): process indices [1,2,3], set l=4

• When j=5 (nums[5]=9): process indices [4,5,6], set l=7

• No overlap! Each index processed exactly once.

This is a clever optimization that transforms what could be O(n × k) worst-case complexity into linear time by avoiding redundant work.

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        res = []
        n = len(nums)
        l = 0
        for j,num in enumerate(nums):
            if num == key:
                left = max(l,j-k) # Start from the next unprocessed index
                right = min(j+k+1, n) # Include j+k
                for i in range(left, right):
                    res.append(i)
                l = right
        return res