Find Shortest Distance to Target Color

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.

Examples:

Input: colors = [1, 1, 2, 1, 3, 2, 2, 3, 3], queries = [[1, 3], [2, 2], [6, 1]]
Output: [3, 0, 3]
Explanation:  The nearest 3 from index 1 is at index 4 (3 steps away). The nearest 2 from index 2 is at index 2 itself (0 steps away). The nearest 1 from index 6 is at index 3 (3 steps away).

Input:  colors = [2, 1, 1, 3, 2, 3, 1, 3, 3] queries = [[1, 2], [5, 3], [7, 1]]
Output: [1, 0, 1]

Explanation:  The nearest 2 from index 1 is at index 0 (1 steps away). The nearest 3 from index 5 is at index 5 itself (0 steps away). The nearest 1 from index 7 is at index 6 (1 steps away).

Approach: To solve the problem, follow the below idea:

The problem can be solved by precomputation the minimum distance to the colors 1, 2 and 3 on the right as well as on the left. For each index i, we can store the minimum distance of color 1, 2 and 3 on the right and on the left and return the minimum between them. We can maintain 2 2D arrays of size N * 3, left[][] and right[][] such that left[i][j] stores the nearest index with color j + 1 on the left of index i(including i) and right[i][j] stores the nearest index with color j + 1 to the right of index i(including index i).

Step-by-step algorithm:

• Initialize 2D array left[][] and right[][] of size N X 3.
• Initialize the first column of left according to the first color in the input array.
• Initialize the last column of right according to the last color in the input array.
• For all indices i > 0,
  • Copy the values from the previous indices, left[i][j] = left[i - 1][j]
  • Update the current color in left, left[i][color[i] - 1] = i;
• For all indices i < N – 1,
  • Copy the values from the next indices, right[i][j] = right[i + 1][j]
  • Update the current color in right, right[i][color[i] - 1] = i;
• For each query, return the minimum of left[] and right[].

Below is the implementation of above approach:

C++

#include <algorithm>
#include <iostream>
#include <vector>

using namespace std;

// Function to find the shortest distance to a target
// color from each query index
vector<int>
shortestDistanceColor(vector<int>& colors,
                      vector<vector<int> >& queries)
{
    // Size of colors
    int N = colors.size();
    int INF = 1000000000;
    // Declare left and right arrays
    vector<vector<int> > right(N, vector<int>(3, 0));
    vector<vector<int> > left(N, vector<int>(3, 0));
    // Initialize the first column of left and last
    // column of right
    for (int j = 0; j < 3; j++) {
        right[N - 1][j] = INF;
        left[0][j] = -INF;
    }

    // Precompute the right array for each index
    for (int i = N - 2; i >= 0; --i) {
        for (int j = 0; j < 3; ++j) {
            right[i][j] = right[i + 1][j];
        }
        right[i][colors[i] - 1] = i;
    }

    // Precompute the left array for each index
    for (int i = 1; i < N; ++i) {
        for (int j = 0; j < 3; ++j) {
            left[i][j] = left[i - 1][j];
        }
        left[i][colors[i - 1] - 1] = i - 1;
    }

    // Vector to store the answer
    vector<int> ans;
    for (vector<int>& q : queries) {
        int i = q[0], c = q[1] - 1;
        // Find the minimum distance using left and
        // right
        int d = min(i - left[i][c], right[i][c] - i);
        // If no nearest color is found
        if (d > N)
            ans.push_back(-1);
        // If nearest color is found, return the
        // distance
        else
            ans.push_back(d);
    }
    return ans;
}

int main()
{
    // Define colors array
    vector<int> colors = { 2, 1, 1, 3, 2, 3, 1, 3, 3 };
    // Define queries array
    vector<vector<int> > queries
        = { { 1, 2 }, { 5, 3 }, { 7, 1 } };
    // Call method to find shortest distances
    vector<int> result
        = shortestDistanceColor(colors, queries);
    // Print result array
    cout << "[";
    for (int i = 0; i < result.size(); ++i) {
        cout << result[i];
        if (i < result.size() - 1)
            cout << ", ";
    }
    cout << "]" << endl;
    return 0;
}
// This code is contributed by Ayush Mishra

Java

import java.util.*;

public class Main {
    // Method to find the shortest distance to a target
    // color from each query index
    public static List<Integer>
    shortestDistanceColor(int[] colors, int[][] queries)
    {
        // Size of colors
        int N = colors.length;
        int INF = 1000000000;
        // Declare left and right arrays
        int[][] right = new int[N][3];
        int[][] left = new int[N][3];
        // Initialize the first column of left and last
        // column of right
        for (int j = 0; j < 3; j++) {
            right[N - 1][j] = INF;
            left[0][j] = -INF;
        }

        // Precompute the right array for each index
        for (int i = N - 2; i >= 0; --i) {
            for (int j = 0; j < 3; ++j) {
                right[i][j] = right[i + 1][j];
            }
            right[i][colors[i] - 1] = i;
        }

        // Precompute the left array for each index
        for (int i = 1; i < N; ++i) {
            for (int j = 0; j < 3; ++j) {
                left[i][j] = left[i - 1][j];
            }
            left[i][colors[i - 1] - 1] = i - 1;
        }

        // Array to store the answer
        List<Integer> ans = new ArrayList<>();
        for (int[] q : queries) {
            int i = q[0], c = q[1] - 1;
            // Find the minimum distance using left and
            // right
            int d
                = Math.min(i - left[i][c], right[i][c] - i);
            // If no nearest color is found
            if (d > N)
                ans.add(-1);
            // If nearest color is found, return the
            // distance
            else
                ans.add(d);
        }
        return ans;
    }

    public static void main(String[] args)
    {
        // Define colors array
        int[] colors = { 2, 1, 1, 3, 2, 3, 1, 3, 3 };
        // Define queries array
        int[][] queries = { { 1, 2 }, { 5, 3 }, { 7, 1 } };
        // Call method to find shortest distances
        List<Integer> result
            = shortestDistanceColor(colors, queries);
        // Print result array
        System.out.println(result.toString());
    }
}

Python

def shortest_distance_color(colors, queries):
    # Size of colors
    N = len(colors)
    INF = 1000000000

    # Declare left and right arrays
    right = [[0] * 3 for _ in range(N)]
    left = [[0] * 3 for _ in range(N)]

    # Initialize the first column of left and last column of right
    for j in range(3):
        right[N - 1][j] = INF
        left[0][j] = -INF

    # Precompute the right array for each index
    for i in range(N - 2, -1, -1):
        for j in range(3):
            right[i][j] = right[i + 1][j]
        right[i][colors[i] - 1] = i

    # Precompute the left array for each index
    for i in range(1, N):
        for j in range(3):
            left[i][j] = left[i - 1][j]
        left[i][colors[i - 1] - 1] = i - 1

    # List to store the answer
    ans = []
    for q in queries:
        i, c = q[0], q[1] - 1
        # Find the minimum distance using left and right
        d = min(i - left[i][c], right[i][c] - i)
        # If no nearest color is found
        if d > N:
            ans.append(-1)
        # If nearest color is found, return the distance
        else:
            ans.append(d)
    return ans


# Define colors array
colors = [2, 1, 1, 3, 2, 3, 1, 3, 3]
# Define queries array
queries = [[1, 2], [5, 3], [7, 1]]
# Call function to find shortest distances
result = shortest_distance_color(colors, queries)
# Print result list
print(result)

JavaScript

// Function to find the shortest distance to a target
// color from each query index
function shortestDistanceColor(colors, queries) {
    // Size of colors
    const N = colors.length;
    const INF = 1000000000;
    // Declare left and right arrays
    let right = new Array(N).fill().map(() => new Array(3).fill(0));
    let left = new Array(N).fill().map(() => new Array(3).fill(0));
    // Initialize the first column of left and last
    // column of right
    for (let j = 0; j < 3; j++) {
        right[N - 1][j] = INF;
        left[0][j] = -INF;
    }

    // Precompute the right array for each index
    for (let i = N - 2; i >= 0; --i) {
        for (let j = 0; j < 3; ++j) {
            right[i][j] = right[i + 1][j];
        }
        right[i][colors[i] - 1] = i;
    }

    // Precompute the left array for each index
    for (let i = 1; i < N; ++i) {
        for (let j = 0; j < 3; ++j) {
            left[i][j] = left[i - 1][j];
        }
        left[i][colors[i - 1] - 1] = i - 1;
    }

    // Array to store the answer
    let ans = [];
    for (let q of queries) {
        let i = q[0], c = q[1] - 1;
        // Find the minimum distance using left and
        // right
        let d = Math.min(i - left[i][c], right[i][c] - i);
        // If no nearest color is found
        if (d > N)
            ans.push(-1);
        // If nearest color is found, return the
        // distance
        else
            ans.push(d);
    }
    return ans;
}

// Define colors array
let colors = [2, 1, 1, 3, 2, 3, 1, 3, 3];
// Define queries array
let queries = [[1, 2], [5, 3], [7, 1]];
// Call method to find shortest distances
let result = shortestDistanceColor(colors, queries);
// Print result array
console.log(result);

Output

[1, 0, 1]

Time Complexity: O(N + Q * log(N)), where N is number of the elements in the input array and and Q is the number of queries.
Auxiliary Space: O(N)