leetcode 1182. 与目标颜色间的最短距离（C++）

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

C++

class Solution {
public:
    vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) 
    {
        vector<int> res;
        int n=colors.size();
        int m=queries.size();
        vector<vector<int>> tmp(4);
        for(int i=0;i<n;i++)
        {
            tmp[colors[i]].push_back(i);
        }
        for(auto it:queries)
        {
            int index=it[0];
            int color=it[1];
            if(0==tmp[color].size())
            {
                res.push_back(-1);
            }
            else
            {
                int left=0;
                int right=tmp[color].size()-1;
                int flag=0;
                while(right-left>1)
                {
                    if(index<=tmp[color][left])
                    {
                        res.push_back(tmp[color][left]-index);
                        flag=1;
                        break;
                    }
                    if(index>=tmp[color][right])
                    {
                        res.push_back(index-tmp[color][right]);
                        flag=1;
                        break;
                    }
                    int mid=(left+right)/2;
                    if(index>=tmp[color][mid])
                    {
                        left=mid;
                    }
                    else
                    {
                        right=mid;
                    }
                }
                if(0==flag)
                {
                    int a1=abs(tmp[color][left]-index);
                    int a2=abs(index-tmp[color][right]);
                    res.push_back(min(a1,a2));
                }
            }
        }
        return res;
    }
};

法2：利用upper_bound

class Solution {
public:
    vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) 
    {
        vector<int> res;
        int n=colors.size();
        int m=queries.size();
        vector<vector<int>> tmp(4);
        for(int i=0;i<n;i++)
        {
            tmp[colors[i]].push_back(i);
        }
        for(auto it:queries)
        {
            int index=it[0];
            int color=it[1];
            if(0==tmp[color].size())
            {
                res.push_back(-1);
            }
            else
            {
                int right=tmp[color].size()-1;
                if(index<=tmp[color][left])
                {
                    res.push_back(tmp[color][left]-index);
                }
                else if(index>=tmp[color][right])
                {
                    res.push_back(index-tmp[color][right]);
                }
                else
                {
                    int idx=upper_bound(tmp[color].begin(),tmp[color].end(),index)-tmp[color].begin();
                    int a1=abs(tmp[color][idx]-index);
                    int a2=abs(index-tmp[color][idx-1]);
                    res.push_back(min(a1,a2));                
                }
            }
        }
        return res;
    }
};