树查询与操作

Question

给定一棵有 N 个节点的树。每个节点都有一个颜色 Ci。我们还获得了 N-1 个查询，并且在每个查询中，我们被告知要销毁之前未销毁的边。每次我们破坏一条边时，我们都必须报告断开连接的节点对的数量。在这里，如果在破坏之前可以使用尚未破坏的边从 i 到达 j，并且如果 Ci=Cj，则称两个节点 i 和 j 断开连接。

如何解决这个问题？我的一些测试用例正在使用 dfs 运行，但有些显示超出时间限制。 请帮忙

我的代码：

#include <bits/stdc++.h>

using namespace std;
#define ll long long int
int par[300002];
int sz[300002];
vector<pair<ll, ll>> arr;
ll n, i, a, b, x, y;

ll find(ll x) {
    if (x != par[x]) {
        par[x] = find(par[x]);
    }
    return par[x];
}

bool same(ll x, ll y) { return find(x) == find(y); }

void merge(ll x, ll y) {
    ll x1 = find(x);
    ll y1 = find(y);

    if (sz[x1] > sz[y1]) {
        par[y] = x;
        sz[x1] += sz[y1];
    } else {
        par[x] = y;
        sz[y1] += sz[x1];
    }
}

int main() {

    cin >> n;
    vector<ll> ans(n + 1);

    for (i = 1; i <= n; i++) {
        par[i] = i;
        sz[i] = 1;
    }

    for (i = 1; i <= n - 1; i++) {
        cin >> a >> b;
        arr.push_back(make_pair(a, b));
    }

    int c[n + 1];
    for (i = 1; i <= n; i++) {
        cin >> c[i];
    }

    ans[n] = (n * (n - 1)) / 2;

    for (i = n - 1; i >= 0; i--) {
        x = arr[i].first;
        y = arr[i].second;

        if (same(x, y)) {
            ans[i] = ans[i + 1];
            continue;
        }

        ll s1 = sz[find(x)];
        ll s2 = sz[find(y)];
        merge(x, y);

        ans[i] = max(ans[i + 1] - (s1 * s2), 0LL);
    }
    int z;

    for (i = 1; i <= n - 1; i++) {
        cin >> z;
        cout << ans[z] << '\n';
    }
}

我的一些测试用例被接受了，但大多数情况下超过了时间限制。我想了很多关于如何优化但无法这样做的问题。请帮助

Answer 1

这是hackerearth上跑步比赛的问题。