Coloring Geysers#11731

Coloring Geysers (Editorial)

Problem Author:	Valerii Kovnatskyi
Prepared by:	Valerii Kovnatskyi
Editorial by:	Valerii Kovnatskyi

Let's consider a vertex $v$ . Let its final color be $c_{v} \neq = - 1$ c_v \ne -1. It is claimed that the eruption of geyser $c_{v}$ c_v must occur after the eruptions of all geysers $u$ for which the following conditions hold:

The distance from $v$ to $u$ does not exceed $d_{u}$ d_u (in other words, $v$ is within the coloring zone of $u$ )
$u \neq = c_{v}$ u \ne c_v

It is not difficult to understand that this set of conditions is not only necessary but also sufficient for the color of vertex $v$ to be exactly $c_{v}$ c_v.

Group 2: $d_{i} = 1$ d_i = 1

For each geyser $u$ , we iterate through all vertices $v$ in the coloring zone. If $c_{v} \neq = v$ c_v \ne v and $c_{v} \neq = - 1$ c_v \ne -1, we draw an edge from $u$ to $c_{v}$ c_v, which will denote the condition that the eruption of geyser $u$ must occur before the eruption of geyser $c_{v}$ c_v. This will work in a number of operations equal to $\sum_{i = 0}^{n} d e g_{v} + 1 \leq 3 n$ \sum_{i=0}^n deg_v + 1 \le 3n (the same number of edges will be in the resulting graph). Such a set of conditions will be identical to the set described above, applied for each vertex.

Now, we will use the topological sorting algorithm to obtain the required order of vertices. The time complexity of the algorithm is $O (n + m)$ O(n + m), therefore the entire solution works in $O (n)$ . An example implementation:

void solve() {
    int n;
    cin >> n;

    vector<int> d(n), c(n);
    cin >> d >> c;
    for (int i = 0; i < n; i++) {
        if (c[i] != -1) c[i] -= 1;
    }

    vector<vector<int>> g(n);
    for (int i = 0; i < n - 1; i++) {
        int v, u;
        cin >> v >> u;
        v -= 1, u -= 1;
        g[v].push_back(u);
        g[u].push_back(v);
    }

    vector<vector<int>> h(n);
    for (int v = 0; v < n; v++) {
        if (c[v] != -1 && c[v] != v) {
            h[v].push_back(c[v]);
        }
        for (int u : g[v]) {
            if (c[u] != -1 && c[u] != v) {
                h[v].push_back(c[u]);
            }
        }
    }

    vector<int> order;
    vector<bool> used(n, false);
    function<void(int)> topsort = [&] (int v) {
        used[v] = true;
        for (int u : h[v]) {
            if (used[u]) continue;
            topsort(u);
        }
        order.push_back(v);
    };
    for (int v = 0; v < n; v++) {
        if (!used[v]) topsort(v);
    }
    reverse(all(order));
    for (int el : order) cout << el + 1 << ' ';
}

For a complete solution, we will use the idea of centroid decomposition. Specifically, we will use it to maintain the progress of the topological sorting algorithm without explicitly constructing the graph in advance.

Let's say we are currently at vertex $v$ . We know that we have edges to those vertices $u$ for which there exists a vertex $w$ such that $c_{w} = u$ c_w = u and the length of the path from $v$ to $w$ does not exceed $d_{v}$ d_v. Let's iterate through the centroid ancestors of vertex $v$ , let the current one be vertex $p r$ . We will check such centroid children $c h$ of vertex $p r$ for which $d i s t (c h, p r) + d i s t (p r, v) \leq d_{v}$ dist(ch, pr) + dist(pr, v) \le d_v, that is, $d i s t (c h, p r) \leq d_{v} - d i s t (p r, v)$ dist(ch, pr) \le d_v - dist(pr, v). It is obvious that from this equality it follows that $d i s t (c h, v) \leq d_{v}$ dist(ch, v) \le d_v, but we can also say that if $p r$ lies on the path between $v$ and $c h$ , then $d i s t (c h, p r) + d i s t (p r, v) = d i s t (c h, v)$ dist(ch, pr) + dist(pr, v) = dist(ch, v). Therefore, if we consider any suitable vertex $w$ , at the moment when we iterate through the centroid children of the least common centroid ancestor of vertices $v$ and $w$ (which, as is known, lies on the path from $v$ to $w$ ), we will definitely check vertex $w$ . Thus, by iterating through the centroid ancestors and centroid children in this way, we will check all the necessary vertices.

We will store two arrays of pairs for each vertex; the first will contain the indices of the centroid ancestors and the distances to them, and the second will contain the indices of the centroid children and the distances to them, with the latter needing to be sorted by distances so that suitable centroid children are stored in some prefix. Both of these arrays can be filled during the construction of the centroid decomposition, either in $O (n \cdot l o g (n)^{2}$ O(n \cdot log(n)^2 time by explicitly sorting the arrays by distances, or in $O (n \cdot l o g (n))$ O(n \cdot log(n)) using breadth-first search.

Now we can perform the topological sorting as described above. We need to apply one optimization to achieve acceptable asymptotic behavior — if we have considered a centroid child $c h$ of centroid ancestor $p r$ , then there is no point in reconsidering it further, since its color has already been used. Therefore, we will store the length of the already considered prefix of the array for each centroid. Since the total size of the stored information is $O (n \cdot l o g (n))$ O(n \cdot log(n)) and we consider each element at most once, the final time complexity will be $O (n \cdot l o g (n))$ O(n \cdot log(n)).

void solve() {
    int n;
    cin >> n;
    vector<int> d(n), c(n);
    cin >> d >> c;
    for (int i = 0; i < n; i++) {
        if (c[i] != -1) c[i] -= 1;
    }
    vector<vector<int>> g(n);
    for (int i = 0; i < n - 1; i++) {
        int v, u;
        cin >> v >> u;
        v -= 1, u -= 1;
        g[v].push_back(u);
        g[u].push_back(v);
    }

    vector<bool> disabled(n, false);
    vector<int> sz(n);
    function<void(int, int)> calc_sz = [&] (int v, int pr) {
        sz[v] = 1;
        for (int u : g[v]) {
            if (disabled[u] || u == pr) continue;
            calc_sz(u, v);
            sz[v] += sz[u];
        }
    };
    function<int(int, int, int)> find_centroid = [&] (int v, int pr, int root) {
        for (int u : g[v]) {
            if (disabled[u] || u == pr) continue;
            if (sz[u] * 2 > sz[root]) return find_centroid(u, v, root);
        }
        return v;
    };
    struct centroid_data {
        int v, dist;
    };
    vector<vector<centroid_data>> centroid_parents(n), centroid_children(n);
    function<void(int)> centroid_decomposition = [&] (int root) {
        calc_sz(root, -1);
        int centroid = find_centroid(root, -1, root);

        struct bfs_data {
            int v, pr, dist;
        };
        queue<bfs_data> q;
        q.push({centroid, -1, 0});
        while (!q.empty()) {
            bfs_data curr = q.front();
            q.pop();
            centroid_children[centroid].push_back({curr.v, curr.dist});
            centroid_parents[curr.v].push_back({centroid, curr.dist});
            for (int u : g[curr.v]) {
                if (disabled[u] || u == curr.pr) continue;
                q.push({u, curr.v, curr.dist + 1});
            }
        }

        disabled[centroid] = true;
        for (int v : g[centroid]) {
            if (!disabled[v]) {
                centroid_decomposition(v);
            }
        }
    };

    centroid_decomposition(0);
    vector<int> checked(n, 0);
    vector<bool> used(n, false);
    vector<int> order;
    function<void(int)> topsort = [&] (int v) {
        used[v] = true;
        for (auto pr : centroid_parents[v]) {
            while (checked[pr.v] < len(centroid_children[pr.v])) {
                centroid_data curr = centroid_children[pr.v][checked[pr.v]];
                if (curr.dist > d[v] - pr.dist) break;
                checked[pr.v]++;
                if (used[c[curr.v]]) continue;
                if (c[curr.v] == -1) continue;
                topsort(c[curr.v]);
            }
        }
        order.push_back(v);
    };

    for (int v = 0; v < n; v++) {
        if (!used[v]) topsort(v);
    }
    reverse(all(order));
    for (int el : order) cout << el + 1 << ' ';
}