Problem statement

https://leetcode.com/problems/shortest-distance-to-target-color/

Solution

Almost identical to 0821. Shortest Distance to a Character, but here we need to calculate not all distances, but given in queries. However we still need to precalculate all distances in advance.

Complexity

Time complexity is O(n*T), where T is number of colors, space complexity is O(n*T) as well.

Code

class Solution:
    def shortestDistanceColor(self, colors, queries):
        def letter_get(letter, dr):
            n = len(colors)
            res, cur = [-1]*n, -n
            for i in range(n)[::dr]:
                if colors[i] == letter: cur = i
                res[i] = abs(i - cur)
            return res

        d, n = defaultdict(list), len(colors)
        for c in range(1, 4):
            d[c] = [min(x, y) for x, y in zip(letter_get(c, 1), letter_get(c, -1))]

        return [d[c][i] if d[c][i] < n else -1 for i, c in queries]