Find All K-Distant Indices in an Array
Problem
Given a 0-indexed integer array nums, an integer key, and an integer k, an index i is k-distant if there exists at least one index j with |i − j| ≤ k and nums[j] == key. Return a sorted list of all k-distant indices.
nums = [3, 4, 9, 1, 3, 9, 5], key = 9, k = 1[1, 2, 3, 4, 5, 6]def find_k_distant_indices(nums, key, k):
n = len(nums)
keep = [False] * n
for j, x in enumerate(nums):
if x == key:
lo = max(0, j - k)
hi = min(n - 1, j + k)
for i in range(lo, hi + 1):
keep[i] = True
return [i for i in range(n) if keep[i]]
function findKDistantIndices(nums, key, k) {
const n = nums.length;
const keep = new Array(n).fill(false);
for (let j = 0; j < n; j++) {
if (nums[j] === key) {
const lo = Math.max(0, j - k);
const hi = Math.min(n - 1, j + k);
for (let i = lo; i <= hi; i++) keep[i] = true;
}
}
const res = [];
for (let i = 0; i < n; i++) if (keep[i]) res.push(i);
return res;
}
class Solution {
public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
int n = nums.length;
boolean[] keep = new boolean[n];
for (int j = 0; j < n; j++) {
if (nums[j] == key) {
int lo = Math.max(0, j - k);
int hi = Math.min(n - 1, j + k);
for (int i = lo; i <= hi; i++) keep[i] = true;
}
}
List<Integer> res = new ArrayList<>();
for (int i = 0; i < n; i++) if (keep[i]) res.add(i);
return res;
}
}
vector<int> findKDistantIndices(vector<int>& nums, int key, int k) {
int n = nums.size();
vector<bool> keep(n, false);
for (int j = 0; j < n; j++) {
if (nums[j] == key) {
int lo = max(0, j - k);
int hi = min(n - 1, j + k);
for (int i = lo; i <= hi; i++) keep[i] = true;
}
}
vector<int> res;
for (int i = 0; i < n; i++) if (keep[i]) res.push_back(i);
return res;
}
Explanation
An index is "k-distant" if it sits within distance k of some position where the value equals key. The simplest reliable approach is to find every key and paint a window of nearby indices as kept.
We use a boolean array keep, all false to start. We scan the array, and whenever nums[j] equals key, we mark the whole window from j - k to j + k as true. We clamp the window with max(0, ...) and min(n - 1, ...) so we never run off the ends.
Windows from different keys may overlap, and that is fine — marking a slot true twice does no harm, so we naturally get the union of all windows. At the end we collect every index whose keep flag is true, already in sorted order since we scan left to right.
Example: nums = [3, 4, 9, 1, 3, 9, 5], key = 9, k = 1. The 9 at index 2 marks [1, 3] and the 9 at index 5 marks [4, 6]. Their union is [1, 2, 3, 4, 5, 6].
Each key can paint up to 2k + 1 cells, so the work is about n · k.