Codeforces - C. Tree Infection

 Problem Link

Problem Summary

A tree with N nodes is provided, where all the nodes are initially healthy. In the end, all nodes must be infected. There are 2 ways to infect a node: inject it with the infection, or let the infection spread. The infection can only spread from one node to another if they have the same parent node. In one second, the infection spreads to all possible nodes, and then one node is injected. Find the minimum number of seconds required for all nodes to be infected.

Algorithm Used

Graph Theory, Greedy, Simulation

Editorial

If  node is injected, then the infection can only spread to nodes within it's group.

If only one node in a group is injected, then the number of seconds the infection takes to spread is the number of nodes in that group. So, to reduce the number of seconds in total, we inject the nodes that are part of the biggest groups first and try to get at least one node from every group injected as quickly as possible, so that the infection spreads sooner.

If one node from all groups has been injected, but not all nodes have been infected yet, then the nodes that have not been infected can be injected.

In my solution, I use two priority queues: one that stores the sizes of the groups that do not have a node injected yet, and another priority queue that stores the sizes of the groups that have nodes that need to be infected. Each time a node is infected, the size of that group goes down by one. Each time the infection spreads and then a node is injected, the number of seconds taken goes up by one (stored in the variable "ans").

Time Complexity

O (N log N)

Code

#include <bits/stdc++.h>

using namespace std;

const int MAX_N = 2 * 1e5 + 1;

vector <vector <int>> adj;

int sz[MAX_N];

void dfs(int s, int p) {

    for (auto i : adj[s]) {

        if (i == p)

            continue;

        sz[s]++;

        dfs(i, s);

    }

}

int main() {

    int T;

    cin >> T;

    while (T--) {

        int n;

        cin >> n;

        adj.clear();

        adj.resize(n+1);

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

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

            int x;

            cin >> x;

            adj[x].push_back(i);

        }

        dfs(1, 0);

        int ans = 0;

        priority_queue <int> inject;

        priority_queue <int> spread_infection;

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

            if (sz[i] != 0)

                inject.push(sz[i]);

        }

        inject.push(1);

        while (!inject.empty() || !spread_infection.empty()) {

            ans++;

            int x = spread_infection.size();

            priority_queue <int> tp;

            while (x--) {

                if (spread_infection.top() - 1 != 0)

                    tp.push(spread_infection.top() - 1);

                spread_infection.pop();

            }

            spread_infection = tp;

            if (!inject.empty()) {

                if (inject.top() - 1 != 0)

                    spread_infection.push(inject.top() - 1);

                inject.pop();

            } else if (!spread_infection.empty()) {

                x = spread_infection.top();

                spread_infection.pop();

                if (x-1 != 0)

                    spread_infection.push(x-1);

            }

        }

        cout << ans << '\n';

    }

}