Problem of The Day: Find All K-Distant Indices in an Array

Problem Statement

leetcode problem link

Brute Force [Accepted]

class Solution:
    def findKDistantIndices(self, nums: List[int], key: int, k: int) -> List[int]:
        key_indices = deque()
        res = set()
        for i, num in enumerate(nums):
            if num == key:
                key_indices.append(i)

        while key_indices:
            curr_index = key_indices.popleft()
            for i in range(len(nums)):
                if abs(i - curr_index) <= k:
                    res.add(i)

        return list(sorted(list(res)))

Editorial

Approach 1: Enumerate

class Solution:
    def findKDistantIndices(
        self, nums: List[int], key: int, k: int
    ) -> List[int]:
        res = []
        n = len(nums)
        # traverse number pairs
        for i in range(n):
            for j in range(n):
                if nums[j] == key and abs(i - j) <= k:
                    res.append(i)
                    break  # early termination to prevent duplicate addition
        return res

Approach 2: One-time Traversal

class Solution:
    def findKDistantIndices(
        self, nums: List[int], key: int, k: int
    ) -> List[int]:
        res = []
        r = 0  # unjudged minimum index
        n = len(nums)
        for j in range(n):
            if nums[j] == key:
                l = max(r, j - k)
                r = min(n - 1, j + k) + 1
                for i in range(l, r):
                    res.append(i)
        return res