The brute force: scan -k..k at each index. More optimal: build suffix array to predict where is the next key, scan and save the last key position.
Approach
use brute-force for easy problems
Complexity
Time complexity: $$O(n^2)$$
Space complexity: $$O(n)$$
Code
// 32ms
fun findKDistantIndices(n: IntArray, key: Int, k: Int) =
n.indices.filter { (max(0, it - k)..min(n.size - 1, it + k)).any { n[it] == key} }
// 1ms
pub fn find_k_distant_indices(n: Vec<i32>, key: i32, k: i32) -> Vec<i32> {
(0..n.len() as i32).filter(|i|
(0.max(i - k)..(i + k + 1).min(n.len() as i32)).any(|j| n[j as usize] == key)).collect()
}
// 0ms
vector<int> findKDistantIndices(vector<int>& n, int key, int k) {
int last = -2 * size(n); vector<int> res{}, next(size(n));
for (int i = size(n) - 1; i >= 0; --i)
next[i] = n[i] == key ? i : (i + 1 < size(n) ? next[i + 1] : 2 * size(n));
for (int i = 0; i < size(n); ++i) {
if (n[i] == key) last = i;
if (next[i] - i <= k || i - last <= k) res.push_back(i);
}
return res;
}
