Find All K-Distant Indices in an Array
EasyPrompt
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.
Example 1:
nums[2] == keyExample 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].
Constraints:
1 <= nums.length <= 10001 <= nums[i] <= 1000keyis an integer from the arraynums.1 <= k <= nums.length
Approaches
4 approaches with complexity analysis and trade-offs.
This approach directly translates the problem definition into code. It iterates through every possible index i in the array and, for each i, performs a full scan of the array to check if it satisfies the k-distant condition.
Algorithm
- Initialize an empty list
result. - Loop with index
ifrom0tonums.length - 1. - Inside the loop, start another loop with index
jfrom0tonums.length - 1. - Check if
nums[j] == keyandMath.abs(i - j) <= k. - If the condition is true, add
itoresultandbreakthe inner loop (since we only need one suchjto exist). - Return
result.
Walkthrough
For each index i from 0 to n-1, we need to determine if it's a k-distant index. To do this, we perform another full scan of the array using an index j. In this inner scan, we look for an index j where nums[j] is equal to the key and the absolute difference |i - j| is less than or equal to k. If such a j is found, we confirm that i is a k-distant index. We then add i to our result list and can immediately stop searching for this i (by breaking the inner loop) and move to the next index i+1. Since we iterate i in increasing order, the resulting list will naturally be sorted.
class Solution { public List<Integer> findKDistantIndices(int[] nums, int key, int k) { List<Integer> result = new ArrayList<>(); int n = nums.length; for (int i = 0; i < n; i++) { for (int j = 0; j < n; j++) { if (nums[j] == key && Math.abs(i - j) <= k) { result.add(i); break; // Found a valid j, move to the next i } } } return result; }}Complexity
Time
O(N^2), where N is the length of the `nums` array. The outer loop runs N times, and for each iteration, the inner loop can run up to N times.
Space
O(N) in the worst case for storing the result list, where N is the number of elements in `nums`. This occurs if all indices are k-distant.
Trade-offs
Pros
Very simple to understand and implement.
Directly follows the problem statement.
Cons
Highly inefficient due to nested loops, leading to a quadratic time complexity.
Likely to result in a 'Time Limit Exceeded' error on platforms with stricter time limits for larger inputs.
This approach improves upon the brute-force method by first identifying all locations of the key. This pre-calculation avoids repeatedly searching the entire array for the key for every single index i.
Algorithm
- Initialize an empty list
keyIndices. - Iterate through
numsfromi = 0ton-1. Ifnums[i] == key, additokeyIndices. - Initialize an empty list
result. - Loop with index
ifrom0ton-1. - Inside, loop through each
jinkeyIndices. - If
Math.abs(i - j) <= k, additoresultandbreakthe inner loop. - Return
result.
Walkthrough
First, we make a single pass through the nums array to find all indices j where nums[j] == key. We store these indices in a separate list, say keyIndices. Then, we iterate through every index i from 0 to n-1 as before. However, for each i, instead of scanning the entire nums array, we only need to check its distance against the pre-calculated indices in keyIndices. We iterate through keyIndices and if we find any index j such that |i - j| <= k, we add i to our result list and break the inner loop to proceed to the next i.
class Solution { public List<Integer> findKDistantIndices(int[] nums, int key, int k) { int n = nums.length; List<Integer> keyIndices = new ArrayList<>(); for (int i = 0; i < n; i++) { if (nums[i] == key) { keyIndices.add(i); } } List<Integer> result = new ArrayList<>(); for (int i = 0; i < n; i++) { for (int j : keyIndices) { if (Math.abs(i - j) <= k) { result.add(i); break; } } } return result; }}Complexity
Time
O(N * M), where N is the length of `nums` and M is the number of times `key` appears in `nums`. The first pass to find keys is O(N), and the main loop is O(N * M). In the worst case, M can be N, leading to O(N^2) complexity.
Space
O(M + N), where M is the number of occurrences of the `key` and N is the length of `nums`. This is for storing `keyIndices` and the `result` list.
Trade-offs
Pros
More efficient than the pure brute-force approach, especially when the
keyis infrequent in the array.
Cons
The worst-case time complexity is still quadratic, which happens when the key appears frequently in the array.
Instead of checking each index i against all key locations, this approach works "outwards" from each occurrence of the key. For each key found, it generates the entire range of k-distant indices around it and adds them to a collection that handles duplicates.
Algorithm
- Initialize a
HashSet<Integer>calledresultSetto store unique indices. - Iterate through
numswith indexifrom0ton-1. - If
nums[i] == key:- Calculate the range
[start, end]wherestart = Math.max(0, i - k)andend = Math.min(nums.length - 1, i + k). - Loop from
p = starttoendand add eachptoresultSet.
- Calculate the range
- Create a new
ArrayListfrom the elements inresultSet. - Sort the list.
- Return the sorted list.
Walkthrough
The logic is that for each index j where nums[j] == key, all indices i in the range [max(0, j - k), min(n-1, j + k)] are k-distant. We iterate through the array, and whenever we find the key, we generate this range of indices. We add all indices from this range into a HashSet. Using a HashSet is convenient as it automatically handles overlaps between ranges generated by different key occurrences, ensuring each index is stored only once. After checking all occurrences of the key, the HashSet contains all unique k-distant indices. Finally, we convert the set to a list and sort it to produce the final output.
class Solution { public List<Integer> findKDistantIndices(int[] nums, int key, int k) { int n = nums.length; Set<Integer> resultSet = new HashSet<>(); for (int i = 0; i < n; i++) { if (nums[i] == key) { int start = Math.max(0, i - k); int end = Math.min(n - 1, i + k); for (int j = start; j <= end; j++) { resultSet.add(j); } } } List<Integer> result = new ArrayList<>(resultSet); Collections.sort(result); return result; }}Complexity
Time
O(N + M*K + R log R), where N is array length, M is key frequency, K is the distance, and R is the result size. The dominant part is often `M*K`. In the worst case (`M` and `K` are `O(N)`), this can be `O(N^2)`.
Space
O(N) to store the `HashSet` and the final result list in the worst case.
Trade-offs
Pros
More efficient than previous approaches if
kis small.The logic of generating ranges from keys is intuitive.
Cons
Requires extra space for the HashSet.
The final sorting step adds overhead (
R log R, where R is the number of results).Can still be slow if
kis large, approaching O(N^2) in the worst case.
This is the most optimal approach, solving the problem in a single pass through the array. It avoids redundant computations and the need for extra data structures like a HashSet by intelligently tracking the range of indices to be added.
Algorithm
- Initialize an empty list
result. - Initialize a pointer
j = 0. This pointer will mark the next index to be potentially added. - Loop with index
ifrom0tonums.length - 1. - If
nums[i] == key:- Calculate the start of the new range:
start = Math.max(j, i - k). - Calculate the end of the new range:
end = Math.min(nums.length - 1, i + k). - Add all indices from
starttoendto theresultlist. - Update
jtoend + 1to prevent adding duplicates and to start the next search from this point.
- Calculate the start of the new range:
- Return
result.
Walkthrough
We iterate through the array nums with an index i. When we find an occurrence of key at nums[i], we know it activates the indices in the range [i-k, i+k]. The key insight is to add indices to our result list in a contiguous, non-overlapping manner. We use a pointer, let's call it j, to keep track of the start of the next block of indices to add. When we find nums[i] == key, the new range of indices to add is from max(j, i - k) to min(n-1, i + k). We add all indices in this calculated range to our result list. Then, we update j to be end + 1. This crucial step ensures that the next time we find a key, we don't re-add indices that are already in our list because the next start will be at least the new j. This process guarantees that we iterate through the indices 0 to n-1 only once in total across all loops, achieving linear time complexity.
class Solution { public List<Integer> findKDistantIndices(int[] nums, int key, int k) { List<Integer> result = new ArrayList<>(); int n = nums.length; // j is the starting point for the next range of indices to add. // It ensures we don't add duplicate indices. int j = 0; for (int i = 0; i < n; i++) { if (nums[i] == key) { // The start of the window for the current key at index i. // We take max with j to avoid re-adding indices from overlapping windows. int start = Math.max(j, i - k); // The end of the window. int end = Math.min(n - 1, i + k); // Add all indices in the calculated [start, end] range. for (int p = start; p <= end; p++) { result.add(p); } // Update j to the next index after the current window. j = end + 1; } } return result; }}Complexity
Time
O(N), where N is the length of the `nums` array. Although there is a nested loop, the pointer `j` is always moved forward, ensuring that each index from `0` to `N-1` is added to the result list at most once. Therefore, the total work is proportional to N.
Space
O(N) in the worst case for storing the result list, where N is the length of `nums`.
Trade-offs
Pros
Highly efficient with a linear time complexity.
Solves the problem in what is effectively a single pass.
Does not require extra space for sets or sorting.
Cons
The logic with the advancing pointer
jmight be slightly less intuitive at first glance compared to more direct approaches.
Solutions
Solution
class Solution {public List<Integer> findKDistantIndices(int[] nums, int key, int k) { int n = nums.length; List<Integer> ans = new ArrayList<>(); for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { if (Math.abs(i - j) <= k && nums[j] == key) { ans.add(i); break; } } } return ans; }}class Solution {public: vector<int> findKDistantIndices(vector<int> &nums, int key, int k) { int n = nums.size(); vector<int> ans; for (int i = 0; i < n; ++i) { for (int j = 0; j < n; ++j) { if (abs(i - j) <= k && nums[j] == key) { ans.push_back(i); break; } } } return ans; }};class Solution: def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]: ans = [] n = len(nums) for i in range(n): if any(abs(i - j) <= k and nums[j] == key for j in range(n)): ans . append(i) return ansVideo walkthrough
Newsletter
One sharp idea, every week
System design and interview prep — short enough to finish.
No spam. Unsubscribe anytime.
Practice
Same difficulty — related problems to reinforce the pattern.