Find All K-Distant Indices in an Array

easy array two pointers

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.

Inputnums = [3, 4, 9, 1, 3, 9, 5], key = 9, k = 1
Output[1, 2, 3, 4, 5, 6]
key occurs at j=2 → window [1, 3]; j=5 → window [4, 6]. Union = {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;
}
Time: O(n · k) Space: O(n)