Description
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:
Input: nums = [3,4,9,1,3,9,5], key = 9, k = 1
Output: [1,2,3,4,5,6]
Explanation: Here, nums[2] == key and nums[5] == key.
- For index 0, |0 - 2| > k and |0 - 5| > k, so there is no j where |0 - j| <= k and nums[j] == key. Thus, 0 is not a k-distant index.
- For index 1, |1 - 2| <= k and nums[2] == key, so 1 is a k-distant index.
- For index 2, |2 - 2| <= k and nums[2] == key, so 2 is a k-distant index.
- For index 3, |3 - 2| <= k and nums[2] == key, so 3 is a k-distant index.
- For index 4, |4 - 5| <= k and nums[5] == key, so 4 is a k-distant index.
- For index 5, |5 - 5| <= k and nums[5] == key, so 5 is a k-distant index.
- For index 6, |6 - 5| <= k and nums[5] == key, so 6 is a k-distant index.
Thus, we return [1,2,3,4,5,6] which is sorted in increasing order.
Example 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] <= 1000key is an integer from the array nums.1 <= k <= nums.length
Solutions
Solution 1: Enumeration
We enumerate the index i in the range [0, n), and for each index i, we enumerate the index j in the range [0, n). If |i - j| ≤ k and nums[j] = key, then i is a K-nearest neighbor index. We add i to the answer array, then break the inner loop and enumerate the next index i.
The time complexity is O(n2), where n is the length of the array nums. The space complexity is O(1).
PythonJavaC++GoTypeScriptRust
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 ans(code-box)
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;
}
}(code-box)
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;
}
};(code-box)
func findKDistantIndices(nums []int, key int, k int) (ans []int) {
for i := range nums {
for j, x := range nums {
if abs(i-j) <= k && x == key {
ans = append(ans, i)
break
}
}
}
return ans
}
func abs(x int) int {
if x < 0 {
return -x
}
return x
}(code-box)
function findKDistantIndices(nums: number[], key: number, k: number): number[] {
const n = nums.length;
const ans: number[] = [];
for (let i = 0; i < n; ++i) {
for (let j = 0; j < n; ++j) {
if (Math.abs(i - j) <= k && nums[j] === key) {
ans.push(i);
break;
}
}
}
return ans;
}(code-box)
impl Solution {
pub fn find_k_distant_indices(nums: Vec<i32>, key: i32, k: i32) -> Vec<i32> {
let n = nums.len();
let mut ans = Vec::new();
for i in 0..n {
for j in 0..n {
if (i as i32 - j as i32).abs() <= k && nums[j] == key {
ans.push(i as i32);
break;
}
}
}
ans
}
}(code-box)
Solution 2: Preprocessing + Binary Search
We can preprocess to get the indices of all elements equal to key, recorded in the array idx. All index elements in the array idx are sorted in ascending order.
Next, we enumerate the index i. For each index i, we can use binary search to find elements in the range [i - k, i + k] in the array idx. If there are elements, then i is a K-nearest neighbor index. We add i to the answer array.
The time complexity is O(n × log n), and the space complexity is O(n). Here, n is the length of the array nums.
PythonJavaC++GoTypeScriptRust
class Solution:
def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
idx = [i for i, x in enumerate(nums) if x == key]
ans = []
for i in range(len(nums)):
l = bisect_left(idx, i - k)
r = bisect_right(idx, i + k) - 1
if l <= r:
ans.append(i)
return ans(code-box)
class Solution {
public List<Integer> findKDistantIndices(int[] nums, int key, int k) {
List<Integer> idx = new ArrayList<>();
for (int i = 0; i < nums.length; i++) {
if (nums[i] == key) {
idx.add(i);
}
}
List<Integer> ans = new ArrayList<>();
for (int i = 0; i < nums.length; ++i) {
int l = Collections.binarySearch(idx, i - k);
int r = Collections.binarySearch(idx, i + k + 1);
l = l < 0 ? -l - 1 : l;
r = r < 0 ? -r - 2 : r - 1;
if (l <= r) {
ans.add(i);
}
}
return ans;
}
}(code-box)
class Solution {
public:
vector<int> findKDistantIndices(vector<int>& nums, int key, int k) {
vector<int> idx;
int n = nums.size();
for (int i = 0; i < n; ++i) {
if (nums[i] == key) {
idx.push_back(i);
}
}
vector<int> ans;
for (int i = 0; i < n; ++i) {
auto it1 = lower_bound(idx.begin(), idx.end(), i - k);
auto it2 = upper_bound(idx.begin(), idx.end(), i + k) - 1;
if (it1 <= it2) {
ans.push_back(i);
}
}
return ans;
}
};(code-box)
func findKDistantIndices(nums []int, key int, k int) (ans []int) {
idx := []int{}
for i, x := range nums {
if x == key {
idx = append(idx, i)
}
}
for i := range nums {
l := sort.SearchInts(idx, i-k)
r := sort.SearchInts(idx, i+k+1) - 1
if l <= r {
ans = append(ans, i)
}
}
return
}(code-box)
function findKDistantIndices(nums: number[], key: number, k: number): number[] {
const n = nums.length;
const idx: number[] = [];
for (let i = 0; i < n; i++) {
if (nums[i] === key) {
idx.push(i);
}
}
const search = (x: number): number => {
let [l, r] = [0, idx.length];
while (l < r) {
const mid = (l + r) >> 1;
if (idx[mid] >= x) {
r = mid;
} else {
l = mid + 1;
}
}
return l;
};
const ans: number[] = [];
for (let i = 0; i < n; ++i) {
const l = search(i - k);
const r = search(i + k + 1) - 1;
if (l <= r) {
ans.push(i);
}
}
return ans;
}(code-box)
impl Solution {
pub fn find_k_distant_indices(nums: Vec<i32>, key: i32, k: i32) -> Vec<i32> {
let n = nums.len();
let mut idx = Vec::new();
for i in 0..n {
if nums[i] == key {
idx.push(i as i32);
}
}
let search = |x: i32| -> usize {
let (mut l, mut r) = (0, idx.len());
while l < r {
let mid = (l + r) >> 1;
if idx[mid] >= x {
r = mid;
} else {
l = mid + 1;
}
}
l
};
let mut ans = Vec::new();
for i in 0..n {
let l = search(i as i32 - k);
let r = search(i as i32 + k + 1) as i32 - 1;
if l as i32 <= r {
ans.push(i as i32);
}
}
ans
}
}(code-box)
Solution 3: Two Pointers
We enumerate the index i, and use a pointer j to point to the smallest index that satisfies j ≥ i - k and nums[j] = key. If j exists and j ≤ i + k, then i is a K-nearest neighbor index. We add i to the answer array.
The time complexity is O(n), where n is the length of the array nums. The space complexity is O(1).
PythonJavaC++GoTypeScriptRust
class Solution:
def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
ans = []
j, n = 0, len(nums)
for i in range(n):
while j < i - k or (j < n and nums[j] != key):
j += 1
if j < n and j <= (i + k):
ans.append(i)
return ans(code-box)
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, j = 0; i < n; ++i) {
while (j < i - k || (j < n && nums[j] != key)) {
++j;
}
if (j < n && j <= i + k) {
ans.add(i);
}
}
return ans;
}
}(code-box)
class Solution {
public:
vector<int> findKDistantIndices(vector<int>& nums, int key, int k) {
int n = nums.size();
vector<int> ans;
for (int i = 0, j = 0; i < n; ++i) {
while (j < i - k || (j < n && nums[j] != key)) {
++j;
}
if (j < n && j <= i + k) {
ans.push_back(i);
}
}
return ans;
}
};(code-box)
func findKDistantIndices(nums []int, key int, k int) (ans []int) {
n := len(nums)
for i, j := 0, 0; i < n; i++ {
for j < i-k || (j < n && nums[j] != key) {
j++
}
if j < n && j <= i+k {
ans = append(ans, i)
}
}
return
}(code-box)
function findKDistantIndices(nums: number[], key: number, k: number): number[] {
const n = nums.length;
const ans: number[] = [];
for (let i = 0, j = 0; i < n; ++i) {
while (j < i - k || (j < n && nums[j] !== key)) {
++j;
}
if (j < n && j <= i + k) {
ans.push(i);
}
}
return ans;
}(code-box)
impl Solution {
pub fn find_k_distant_indices(nums: Vec<i32>, key: i32, k: i32) -> Vec<i32> {
let n = nums.len();
let mut ans = Vec::new();
let mut j = 0;
for i in 0..n {
while j < i.saturating_sub(k as usize) || (j < n && nums[j] != key) {
j += 1;
}
if j < n && j <= i + k as usize {
ans.push(i as i32);
}
}
ans
}
}(code-box)
