2200. Find All K-Distant Indices in an Array

SourceWeekly Contest 284 Q1DifficultyEasyRating1266

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 <= 1000
1 <= nums[i] <= 1000
key is an integer from the array nums.
1 <= k <= nums.length

Solutions

Solution 1: Enumeration

Thinking

We need every index within distance \(k\) of some occurrence of \(key\). Checking a window of width \(2k\) around each \(i\), or scanning the whole array for a valid \(j\), is \(O(n^2)\). With \(n \le 10^3\) this bound is acceptable.

Only existence of a \(key\) in \([i-k, i+k]\) matters. Fix \(i\), scan \(j\), and as soon as \(|i-j| \le k\) and \(nums[j] = key\), record \(i\) and leave the inner loop.

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| \leq 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(n^2)\), where \(n\) is the length of the array \(nums\). The space complexity is \(O(1)\).

Python3JavaC++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

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;
    }
}

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;
    }
};

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
}

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;
}

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
    }
}

Solution 2: Preprocessing + Binary Search

Thinking

Solution 1 rescan the array for every \(i\) and never reuses the known positions of \(key\). Those positions can be collected once and queried in logarithmic time.

Store every index of \(key\) in a sorted list \(idx\). For each \(i\), binary-search \(idx\) for a value in \([i-k, i+k]\) via \(\textit{bisect\_left}\) and \(\textit{bisect\_right}\). If \(l \le r\), then \(i\) is valid. The time becomes \(O(n \log n)\).

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 \times \log n)\), and the space complexity is \(O(n)\). Here, \(n\) is the length of the array \(nums\).

Python3JavaC++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

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;
    }
}

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;
    }
};

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
}

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;
}

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
    }
}

Solution 3: Two Pointers

Thinking

Solution 2 already avoids rescanning \(key\), but still stores an index array and binary-searches twice per \(i\). As \(i\) increases, the leftmost \(j\) with \(j \ge i-k\) and \(nums[j] = key\) only moves right.

Advance a pointer \(j\) monotonically: increment while \(j < i-k\) or the current cell is not \(key\). If afterwards \(j \le i+k\), then \(i\) is a K-distant index. One pass, \(O(1)\) extra space.

We enumerate the index \(i\), and use a pointer \(j\) to point to the smallest index that satisfies \(j \geq i - k\) and \(nums[j] = key\). If \(j\) exists and \(j \leq 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)\).

Python3JavaC++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

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;
    }
}

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;
    }
};

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
}

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;
}

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
    }
}

2200. Find All K-Distant Indices in an Array

Description

Solutions

Solution 1: Enumeration

Solution 2: Preprocessing + Binary Search

Solution 3: Two Pointers

Comments