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