LeetCode-Python-1182. 与目标颜色间的最短距离

给你一个数组 colors，里面有  1、2、 3 三种颜色。

我们需要在 colors 上进行一些查询操作 queries，其中每个待查项都由两个整数 i 和 c 组成。

现在请你帮忙设计一个算法，查找从索引 i 到具有目标颜色 c 的元素之间的最短距离。

如果不存在解决方案，请返回 -1。

 

示例 1：

输入：colors = [1,1,2,1,3,2,2,3,3], queries = [[1,3],[2,2],[6,1]]
输出：[3,0,3]
解释： 
距离索引 1 最近的颜色 3 位于索引 4（距离为 3）。
距离索引 2 最近的颜色 2 就是它自己（距离为 0）。
距离索引 6 最近的颜色 1 位于索引 3（距离为 3）。

示例 2：

输入：colors = [1,2], queries = [[0,3]]
输出：[-1]
解释：colors 中没有颜色 3。

 

提示：

• 1 <= colors.length <= 5*10^4
• 1 <= colors[i] <= 3
• 1 <= queries.length <= 5*10^4
• queries[i].length == 2
• 0 <= queries[i][0] < colors.length
• 1 <= queries[i][1] <= 3

思路：

先建立三个数组分别记录所有1， 2， 3的坐标，

然后根据每次查询，在对应的数组里用二分查最近的一个或者两个数的下标，返回差值较小的那个差值即可。

class Solution(object):
    def shortestDistanceColor(self, colors, queries):
        """
        :type colors: List[int]
        :type queries: List[List[int]]
        :rtype: List[int]
        """
        one = []
        two = []
        three = []
        for i, color in enumerate(colors): #预处理
            if color == 1:
                one.append(i)
            elif color == 2:
                two.append(i)
            else:
                three.append(i)
                
        def find(i, c): #三个数组都满足升序，所以可以直接用二分
            if c == 1:
                l = one
            elif c == 2:
                l = two
            else:
                l = three
            
            if len(l) == 0:
                return -1
            
            index = bisect.bisect(l, i) #调库找到插入位置
            # print l, index, i
            if index == 0:
                return l[0] - i
            elif index == len(l):
                return i - l[-1]
            else:
                return min(l[index] - i, i - l[index - 1])
            
        res = []
        for i, c in queries:
            if colors[i] == c:
                res.append(0)
            else:
                res.append(find(i, c))
        
        return res