Codeforces - F. Vlad and Unfinished Business

 Problem Link

Problem Summary

You live at the start house \(X\), and you need to reach house \(Y\) only after completing \(K\) tasks, which are located at some houses. A task is completed once you visit the house where it's located. The houses are numbered from \(1\) to \(N\), and are connected by \(N-1\) roads. Traversing a road takes \(1\) second. Find the minimum time required to complete all the tasks and end at house \(Y\).

Algorithm Used

Tree, DFS

Editorial

The layout of the houses and roads can be modelled as a tree, where the houses are the nodes, the roads are the edges. There is only one path from one node to another, which is why the graph is a tree.

Let's define a few variables;
  1. \(S\) - the node that we start from
  2. \(E\) - the node that we must end at
  3. \(T\) - a node that contains a task
In the picture above, we see that an edge is only traversed if it leads to a \(T\) or and \(E\) node. More specifically, if an edge is on the simple path from the \(S\) node to the \(E\) node, then the edge is only traversed once. Otherwise, if the edge leads to a T node and is not on the simple path from the S node to the T node, it is traversed twice.

Keeping these observations in mind, the algorithm goes as follows: root the tree at the S node, and run a DFS on the tree and store a bool value for each node that indicates whether or not a T node is in the subtree of the current node or not. Run another DFS on the tree and store another bool value that indicates whether or not the E node is in the subtree of the current node or not. These bool values can be stored in an array, where the index of the node is the index to the array.

Run a third DFS to calculate the final answer. If the E node is in the subtree of the current node, then the path is only traversed once, so only 1 second is added to the answer. Otherwise if a T node is in the subtree of the current node and the E node is not, then this path is not the simple path from the S node to the E node, then this edge must be traversed twice, so 2 seconds are added to the answer. Otherwise if both a T node and the E node are not in the subtree of the current node, then there's no point in traversing the path.

Code

#include <bits/stdc++.h>
using namespace std;
const int MAX_N = 2 * 1e5 + 1;
bool task[MAX_N], tsks_in_children[MAX_N], end_in_children[MAX_N];
vector <vector <int>> adj;
int n, k, s, e;

bool dfs(int s, int p) {
    bool a = false;
    if (task[s])
        a = true;
    for (auto i : adj[s]) {
        if (i == p)
            continue;
        a |= dfs(i, s);
    }
    tsks_in_children[s] = a;
    return a;
}

bool dfs_ed(int s, int p) {
    bool a = false;
    if (s == e) {
        a = true;
        end_in_children[s] = true;
        return true;
    }
    for (auto i : adj[s]) {
        if (i == p)
            continue;
        a |= dfs_ed(i, s);
    }
    end_in_children[s] = a;
    return a;
}

int ans(int s, int p) {
    int a = 0;
    for (auto i : adj[s]) {
        if (i == p)
            continue;
        if (end_in_children[i]) {
            a++;
            a += ans(i, s);
        } else if (tsks_in_children[i]) {
            a += 2;
            a += ans(i, s);
        }
    }
    return a;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int T;
    cin >> T;
    while (T--) {
        cin >> n >> k >> s >> e;
        adj.clear();
        adj.resize(n+1);
        fill(task, task + n + 1, false);
        fill(tsks_in_children, tsks_in_children + n + 1, false);
        fill(end_in_children, end_in_children + n + 1, false);
        for (int i = 0; i < k; i++) {
            int x;
            cin >> x;
            task[x] = true;
        }
        for (int i = 0; i < n-1; i++) {
            int a, b;
            cin >> a >> b;
            adj[a].push_back(b);
            adj[b].push_back(a);
        }
        dfs(s, 0);
        dfs_ed(s, 0);
        cout << ans(s, 0) << '\n';
    }
}

Comments

Post a Comment