1182. Shortest Distance to Target Color

You are given an array colors, in which there are three colors: 1, 2 and 3.

You are also given some queries. Each query consists of two integers i and c, return the shortest distance between the given index i and the target color c. If there is no solution return -1.

• binary search solution
• Time : O(n + mlogn)
• Space : O(n)

class Solution {
public:
    vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) {
        vector<vector<int>> order(4);
        for(int i = 0; i < colors.size(); i++) {
            order[colors[i]].push_back(i);
        }
        vector<int> res(queries.size(), -1);
        for(int i = 0; i < queries.size(); i++) {
            int target = queries[i][1], index = queries[i][0];
            if(colors[index] == target) {
                res[i] = 0;
            } else if(!order[target].empty()) {
                auto it = lower_bound(order[target].begin(), order[target].end(), index);
                int ans = INT_MAX;
                if(it != order[target].end()) {
                    ans = min(ans, *it - index);
                }
                if(it != order[target].begin()) {
                    ans = min(ans, index - *(--it));
                }
                res[i] = ans;
            }
        }
        return res;
    }
};

• dynamic programming solution
• Time : O(n + m)
• Space : O(n)

class Solution {
public:
    vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) {
        int n = colors.size();
        vector<vector<int>> l(4, vector<int>(n, n + 1));
        vector<vector<int>> r(4, vector<int>(n, n + 1));
        
        for(int c = 1; c <= 3; c++) {
            r[colors[n-1]][n-1] = l[colors[0]][0] = 0; //init
            for(int i = 1; i < n; i++) {
                l[c][i] = colors[i] == c ? 0 : l[c][i-1] + 1;
                r[c][n-i-1] = colors[n-i-1] == c ? 0 : r[c][n-i] + 1;
            }
        }
        
        vector<int> res;
        for(auto q : queries) {
            int ans = min(l[q[1]][q[0]], r[q[1]][q[0]]);
            if(ans >= n + 1) ans = -1;
            res.push_back(ans);
        }
        return res;
    }
};