aboutsummaryrefslogtreecommitdiff
path: root/find_k_distant_indices.py
diff options
context:
space:
mode:
Diffstat (limited to 'find_k_distant_indices.py')
-rw-r--r--find_k_distant_indices.py34
1 files changed, 34 insertions, 0 deletions
diff --git a/find_k_distant_indices.py b/find_k_distant_indices.py
new file mode 100644
index 0000000..11ddfc6
--- /dev/null
+++ b/find_k_distant_indices.py
@@ -0,0 +1,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))