Skip to content

2200. Find All K-Distant Indices in an Array 👍

  • Time: $O(n)$
  • Space: $O(|\texttt{ans}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution {
 public:
  vector<int> findKDistantIndices(vector<int>& nums, int key, int k) {
    const int n = nums.size();
    vector<int> ans;

    for (int i = 0, j = 0; i < n; ++i) {
      // the first index j s.t. nums[j] == key and j >= i - k
      while (j < n && (nums[j] != key || j < i - k))
        ++j;
      if (j == n)
        break;
      if (abs(i - j) <= k)
        ans.push_back(i);
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution {
  public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
    final int n = nums.length;
    List<Integer> ans = new ArrayList<>();

    for (int i = 0, j = 0; i < n; ++i) {
      // the first index j s.t. nums[j] == key and j >= i - k
      while (j < n && (nums[j] != key || j < i - k))
        ++j;
      if (j == n)
        break;
      if (Math.abs(i - j) <= k)
        ans.add(i);
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution:
  def findKDistantIndices(self, nums: list[int], key: int, k: int) -> list[int]:
    n = len(nums)
    ans = []

    j = 0
    for i in range(n):
      # the first index j s.t. nums[j] == key and j >= i - k
      while j < n and (nums[j] != key or j < i - k):
        j += 1
      if j == n:
        break
      if abs(i - j) <= k:
        ans.append(i)

    return ans