aboutsummaryrefslogtreecommitdiff
path: root/find_k_distant_indices.py
blob: 11ddfc65fcc0e1f9adf4d97708708d6355a54f19 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
# 2200. Find All K-Distant Indices in an Array
import bisect

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


class Solution(object):
    def findKDistantIndices(self, nums, key, k):
        """
        :type nums: List[int]
        :type key: int
        :type k: int
        :rtype: List[int]
        """
        keys = [i for (i, j) in enumerate(nums) if j == key]
        ans = set(keys)
        for i in range(len(nums)):
            idx = bisect.bisect_left(keys, i)
            if 0 <= idx < len(keys) and abs(i - keys[idx]) <= k:
                ans.add(i)
            if 0 < idx <= len(keys) and abs(i - keys[idx - 1]) <= k:
                ans.add(i)
        return sorted(list(ans))


if __name__ == "__main__":
    obj = Solution()
    print(obj.findKDistantIndices(nums=[3, 4, 9, 1, 3, 9, 5], key=9, k=1))
    print(obj.findKDistantIndices(nums=[2, 2, 2, 2, 2], key=2, k=2))