USACO 2021 December Contest, Silver Problem 2 - Connecting Two Barns

Problem Link

Algorithm Used: Connected Components (DFS), binary search, two pointer
Time Complexity: O(T(N + M))

In this problem, T test cases are given. In each test case, there are N fields and M roads, which connect two fields. FJ can build at most 2 roads. The cost of building a new road between field i and field j is (i - j)². There is a barn located in field 1 and another in field N. Find the minimum cost to get from the first to the second barn.

This problem can be modelled as a graph problem, where the fields are the node and the roads are the edges. Roads are bidirectional and have a weight of 0. Only the new roads will be weighted, depending on the nodes they connect.

There are three possible cases: FJ may build 0, 1 or 2 roads. The only case where FJ can build no roads and still reach the second barn is when the two barns are already connected by existing roads, which means that they are in the same component. 

The component containing the first barn will be called the source component and the component containing the second barn will be called the destination component.

If only one road will allow FJ to reach the second barn with the minimum cost, then exactly two components will need to be connected: the source component and the destination component. Otherwise, two roads will be needed, which will connect exactly three components: the source component, another component, and the destination component. In this case, the second component of the three described will act as the intermediate component, to which the other two components will connect. 

The only way to figure out whether one or two roads will be needed is try out the different possibilities and calculate the cost for each one.

To do this, first the components need to be found, which can be done using a simple DFS. If the first and second barn are in the same component, then print 0 as the answer for that test case.

Make a list which contains, for each component, the nodes that it contains. Make sure that the nodes are in sorted order for each list. 

Then comes the tricky bit. For each node in the source component, we could loop over all the other nodes, calculate the cost of building a road between them and store the minimum cost of building a road between the source component and all other components. The minimum cost for the destination component can be calculated similarly. However, this approach would exceed time limit.

Instead, we use the two pointer method. Iterate over each of the nodes. The ith node will act as the intermediate node, to which the source and destination components will connect. Maintain a pointer that points to the first node in the source component that is at least as large as the ith node. As the intermediate node increases, the pointer moves along.

Since the nodes are in sorted order and i is always increasing, two pointer method can be used and it ensures that the nodes are always as close to the ith node as possible. Because of this, the cost will also be as small as possible. 

Calculate the cost of adding this edge. Store the minimum of all edges from the ith node to the source component. The costs for the destination node can be calculated similarly.

Finally, loop over all nodes and add the costs of connecting the source and destination components to the ith node. The answer is the minimum of all the sums.

Code

#include <bits/stdc++.h>

using namespace std;

const int MAX_N = 1e5 + 1;

int comp[MAX_N];

bool visited[MAX_N];

vector <vector <int>> adj;

void dfs(int s, int c) {

    if (visited[s])

        return;

    visited[s] = true;

    comp[s] = c;

    for (int i : adj[s])

        dfs(i, c);

}

int main()

{

    ios_base::sync_with_stdio(0);

    cin.tie(0);

    cout.tie(0);

    int T;

    cin >> T;

    for (int t = 0; t < T; t++) {

        int n, m;

        cin >> n >> m;

        adj.clear();

        adj.resize(n+1);

        fill(comp, comp + n + 1, 0);

        fill(visited, visited + n + 1, false);

        for (int i = 0; i < m; i++) {

            int a, b;

            cin >> a >> b;

            adj[a].push_back(b);

            adj[b].push_back(a);

        }

        int components = 1;

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

            if (!visited[i]) {

                dfs(i, components);

                components++;

            }

        }

        int src_comp = comp[1], dest_comp = comp[n];

        if (src_comp == dest_comp) {

            cout << 0 << '\n';

            continue;

        }

        vector <vector <int>> comp_list(components);

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

            comp_list[comp[i]].push_back(i);

        }

        vector <long long> min_cost_src(components, 1e11);

        vector <long long> min_cost_dest(components, 1e11);

        int src_p = 0, dest_p = 0;

        /*

         * find the minimum cost connecting the src component and the dest

         * component to node i.

         */

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

            int comp_i = comp[i];            

            // first connect the source component

            while (src_p < comp_list[src_comp].size()) {

                // make sure to not overshoot bounds

                int node = comp_list[src_comp][src_p];

                long long cost = pow(abs(i - node), 2);

                min_cost_src[comp_i] = min(min_cost_src[comp_i], cost);

                // move pointer

                if (node < i)

                    src_p++;

                else

                    break;

            }

            if (src_p > 0)

                src_p--;

            //now connect the destiination component

            while (dest_p < comp_list[dest_comp].size()) {

                // make sure not to overshoot bounds

                int node = comp_list[dest_comp][dest_p];

                long long cost = pow(abs(i - node), 2);

                min_cost_dest[comp_i] = min(min_cost_dest[comp_i], cost);

                // move pointer

                if (node < i)

                    dest_p++;

                else

                    break;

            }

            if (dest_p > 0)

                dest_p--;

        }

        long long ans = LONG_LONG_MAX;

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

            ans = min(ans, min_cost_src[i] + min_cost_dest[i]);

        }

        cout << ans << '\n';

    }

    return 0;

}