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

English Version

题目描述

给你一个数组 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

解法

二分查找。

先用哈希表记录每种颜色的索引位置。然后遍历 queries，如果当前 color 不在哈希表中，说明不存在解决方案，此时此 query 对应的结果元素是 -1。否则，在哈希表中取出当前 color 对应的索引列表，二分查找即可。

Python3

class Solution:
    def shortestDistanceColor(self, colors: List[int], queries: List[List[int]]) -> List[int]:
        color_indexes = defaultdict(list)
        for i, c in enumerate(colors):
            color_indexes[c].append(i)
        res = []
        for i, c in queries:
            if c not in color_indexes:
                res.append(-1)
            else:
                t = color_indexes[c]
                left, right = 0, len(t) - 1
                while left < right:
                    mid = (left + right) >> 1
                    if t[mid] >= i:
                        right = mid
                    else:
                        left = mid + 1
                val = abs(t[left] - i)
                if left > 0:
                    val = min(val, abs(t[left - 1] - i))
                if left < len(t) - 1:
                    val = min(val, abs(t[left + 1] - i))
                res.append(val)
        return res

Java

class Solution {
    public List<Integer> shortestDistanceColor(int[] colors, int[][] queries) {
        Map<Integer, List<Integer>> colorIndexes = new HashMap<>();
        for (int i = 0; i < colors.length; ++i) {
            int c = colors[i];
            colorIndexes.computeIfAbsent(c, k -> new ArrayList<>()).add(i);
        }
        List<Integer> res = new ArrayList<>();
        for (int[] query : queries) {
            int i = query[0], c = query[1];
            if (!colorIndexes.containsKey(c)) {
                res.add(-1);
                continue;
            }
            List<Integer> t = colorIndexes.get(c);
            int left = 0, right = t.size() - 1;
            while (left < right) {
                int mid = (left + right) >> 1;
                if (t.get(mid) >= i) {
                    right = mid;
                } else {
                    left = mid + 1;
                }
            }
            int val = Math.abs(t.get(left) - i);
            if (left > 0) {
                val = Math.min(val, Math.abs(t.get(left - 1) - i));
            }
            if (left < t.size() - 1) {
                val = Math.min(val, Math.abs(t.get(left + 1) - i));
            }
            res.add(val);
        }
        return res;
    }
}

...