USACO 2021 January Contest, Gold Problem 2 - Telephone

 Problem Link

Time Complexity: O(Nlog(KlogK))
Algorithm used: Dijkstra's algorithm, Binary Search

This is where I understood how to solve this problem in time.

In this problem, the breeds of N cows are given, followed by a K by K matrix, describing whether or not a cow of breed i is willing to transmit a message to cow of breed j. The cost of sending a message from cow i to cow j is |i - j|. Find the minimum possible cost of transmitting a message from cow 1 to cow N.

This problem can be represented as a graph problem, where the cows are the nodes and edges are drawn between the nodes depending on the compatibility of cow breeds, described in the matrix. While reading in the matrix, also create an adjacency list for every breed.

Intuitively, we could run a standard Dijkstra's algorithm starting from cow 1 and loop over the adjacency list for every pass of the algorithm. Although this implementation would always produce the correct answer, the program would exceed time limit, since there will be N² edges to examine for every pass of the algorithm (time complexity: O(Nlog(N²)). To solve this problem, we just need to examine fewer edges.

When node a is being processed, choose the two adjacent nodes to a from each breed's adjacency list and reduce these nodes' distances. For example, let's say that breed 1 has cows 2, 5, 6 and node a = 3. Here, the two adjacent nodes are 2 and 5. Nodes 2 and 5's distances will get reduced according to standard Dijkstra's algorithm's logic and then added to the queue to be processed. This way, we only examine at most 2K edges for every node being processed. We examine the two adjacent nodes to a because those are the two possible ways a message can be transmitted (forwards and backwards). Examining any other nodes would be redundant since that would only increase the total cost.

I use binary search to find the two adjacent nodes of node a in each breed's adjacency list.

There is one exception to this modified Dijkstra's algorithm: if node 1 and node N are part of the same breed, then the answer may come out to be -1. If node N's breed is not compatible with it's own breed, then only the smallest node in node N's breed's adjacency list will get visited. To fix this, check: if the current node being examined has the same breed as cow N, then answer could be i - 1 + |i - N|, as the cost to travel from cow 1 to cow i is i - 1 and the cost to travel from cow i to cow N is |i - N|.

Code

#include <bits/stdc++.h>

using namespace std;

const int MAX_K = 51;

vector <int> dis;

priority_queue <pair <int, int>> q;

vector <bool> processed;

//adjacency list

vector <vector <int>> adj_list(MAX_K);

//adjacency matrix

bool adj_matrix[MAX_K][MAX_K];

void update_dijkstra(int a, int x){

    if (dis[a] + abs(x - a) < dis[x]){

        dis[x] = dis[a] + abs(x - a);

        q.push({-dis[x], x});

    }

}

int main()

{

    int n, k;

    cin >> n >> k;

    vector <int> breeds(n+1);

    for (int i = 1; i <= n; i++){

        cin >> breeds[i];

        adj_list[breeds[i]].push_back(i);

    }

    for (int i = 1; i <= k; i++){

        string s;

        cin >> s;

        for (int j = 0; j < k; j++)

            adj_matrix[i][j+1] = s[j] - '0';

    }

    q.push({0, 1});

    dis.resize(n+1, INT_MAX);

    processed.resize(n+1, false);

    dis[1] = 0;

    int ans = INT_MAX;

    while (!q.empty()){

        int a = q.top().second;

        q.pop();

        if (processed[a])

            continue;

        processed[a] = true;

        //check for exception

        if (adj_matrix[breeds[a]][breeds[n]] || a == n)

            ans = min(ans, dis[a] + (n - a));

        //loop over 2K breeds

        for (int i = 1; i <= k; i++){

            if (adj_matrix[breeds[a]][i]){

                //find closest left and right node for breed i

                //upper_bound uses binary search

                auto right = upper_bound(adj_list[i].begin(), adj_list[i].end(), a);

                //find closest right node

                if (right != adj_list[i].end())

                    update_dijkstra(a, *right);

                //if breed of cow a = i, then cow a will be included in adj_list[i]. In this case,

                //skip node i and take the node to the left of it.

                if (breeds[a] != i && (right - adj_list[i].begin()) >= 1)

                    update_dijkstra(a, *(right - 1));

                else if (breeds[a] == i && (right - adj_list[i].begin()) >= 2)

                    update_dijkstra(a, *(right - 2));

            }

        }

    }

    ans = min(ans, dis[n]);

    if (ans == INT_MAX)

        cout << -1 << '\n';

    else

        cout << ans << '\n';

    return 0;

}