Shortest Distance to Target Color

medium array prefix precompute

Problem

You are given an array colors where each value is 1, 2, or 3, and a list of queries. Each query [index, color] asks for the shortest distance from index to any element equal to color. Return that distance, or -1 if the color never appears.

Inputcolors = [1, 1, 2, 1, 3, 2, 2, 3, 3], query = [1, 3]

Output3

From index 1, the nearest color 3 sits at index 4, so the distance is |4 − 1| = 3.

colors (comma-separated, values 1/2/3)

index

color

def shortest_distance_color(colors, queries):
    n = len(colors)
    INF = float('inf')
    # nearest[c][i] = distance from i to nearest color c
    nearest = [[INF] * n for _ in range(4)]
    for c in range(1, 4):
        # left-to-right pass
        for i in range(n):
            if colors[i] == c:
                nearest[c][i] = 0
            elif i > 0:
                nearest[c][i] = nearest[c][i - 1] + 1
        # right-to-left pass
        for i in range(n - 2, -1, -1):
            nearest[c][i] = min(nearest[c][i], nearest[c][i + 1] + 1)
    res = []
    for index, color in queries:
        d = nearest[color][index]
        res.append(-1 if d == INF else d)
    return res

function shortestDistanceColor(colors, queries) {
  const n = colors.length, INF = Infinity;
  const nearest = [0, 1, 2, 3].map(() => new Array(n).fill(INF));
  for (let c = 1; c <= 3; c++) {
    for (let i = 0; i < n; i++) {
      if (colors[i] === c) nearest[c][i] = 0;
      else if (i > 0) nearest[c][i] = nearest[c][i - 1] + 1;
    }
    for (let i = n - 2; i >= 0; i--) {
      nearest[c][i] = Math.min(nearest[c][i], nearest[c][i + 1] + 1);
    }
  }
  return queries.map(([index, color]) => {
    const d = nearest[color][index];
    return d === INF ? -1 : d;
  });
}

class Solution {
    public List<Integer> shortestDistanceColor(int[] colors, int[][] queries) {
        int n = colors.length, INF = Integer.MAX_VALUE / 2;
        int[][] nearest = new int[4][n];
        for (int c = 1; c <= 3; c++) {
            for (int i = 0; i < n; i++) {
                if (colors[i] == c) nearest[c][i] = 0;
                else nearest[c][i] = (i > 0) ? nearest[c][i - 1] + 1 : INF;
            }
            for (int i = n - 2; i >= 0; i--) {
                nearest[c][i] = Math.min(nearest[c][i], nearest[c][i + 1] + 1);
            }
        }
        List<Integer> res = new ArrayList<>();
        for (int[] q : queries) {
            int d = nearest[q[1]][q[0]];
            res.add(d >= INF ? -1 : d);
        }
        return res;
    }
}

vector<int> shortestDistanceColor(vector<int>& colors, vector<vector<int>>& queries) {
    int n = colors.size(), INF = INT_MAX / 2;
    vector<vector<int>> nearest(4, vector<int>(n, INF));
    for (int c = 1; c <= 3; c++) {
        for (int i = 0; i < n; i++) {
            if (colors[i] == c) nearest[c][i] = 0;
            else if (i > 0) nearest[c][i] = nearest[c][i - 1] + 1;
        }
        for (int i = n - 2; i >= 0; i--) {
            nearest[c][i] = min(nearest[c][i], nearest[c][i + 1] + 1);
        }
    }
    vector<int> res;
    for (auto& q : queries) {
        int d = nearest[q[1]][q[0]];
        res.push_back(d >= INF ? -1 : d);
    }
    return res;
}

Explanation

There can be many queries, so instead of searching the array fresh each time, we precompute the nearest distance to every color ahead of time. Then each query is just a constant-time lookup.

For each color c in {1, 2, 3} we build a distance row nearest[c] using two passes. In the left-to-right pass, a cell equal to c gets distance 0; otherwise it inherits its left neighbor's distance plus one (the nearest c on the left). In the right-to-left pass we take the minimum with the right neighbor's distance plus one, capturing the nearest c on the right.

After both passes, nearest[c][i] holds the true shortest distance from index i to any cell of color c. A query [index, color] just reads nearest[color][index], returning -1 if that value is still infinity (the color never appears).

Example: colors = [1,1,2,1,3,2,2,3,3], query [1, 3]. The nearest color 3 to index 1 sits at index 4, so the distance is |4 - 1| = 3.

Precomputing costs O(n) per color (a constant three colors), and every query is then O(1), giving overall O(n + q) time.

Time: O(n + q) Space: O(n)