TLE on test18 求调
查看原帖
TLE on test18 求调
307211
ChickyHas楼主2023/11/17 11:46
#include <iostream>
#include <vector>

using namespace std;
using ll = long long;

const int N = 1e5 + 1;

int n, a[N], hson[N], siz[N], cnt[N], mx;
vector<int> e[N];
ll ans[N], sum;

void Dfs1(int x, int fa) {
  for (int i : e[x]) {
    if (i == fa) {
      continue;
    }
    Dfs1(i, x);
    siz[x] += siz[i];
    if (siz[i] > siz[hson[x]]) {
      hson[x] = i;
    }
  }
}

void Dfs3(int x, int fa) {
  cnt[a[x]]++;
  if (cnt[a[x]] > mx) {
    mx = cnt[a[x]], sum = a[x];
  } else if (cnt[a[x]] == mx) {
    sum += a[x];
  }
  for (int i : e[x]) {
    if (i == fa || i == hson[x]) {
      continue;
    }
    Dfs3(i, x);
  }
}

void Init(int x, int fa) {
  cnt[a[x]]--;
  for (int i : e[x]) {
    if (i == fa) {
      continue;
    }
    Init(i, x);
  }
}

void Dfs2(int x, int fa) {
  for (int i : e[x]) {
    if (i == fa || i == hson[x]) {
      continue;
    }
    Dfs2(i, x);
    Init(i, x), sum = mx = 0;
  }
  if (hson[x]) {
    Dfs2(hson[x], x);
  }
  Dfs3(x, fa);
  ans[x] = sum;
}

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0), cout.tie(0);
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    cin >> a[i];
  }
  for (int i = 1, x, y; i < n; ++i) {
    cin >> x >> y;
    e[x].emplace_back(y), e[y].emplace_back(x);
  }
  Dfs1(1, 0);
  Dfs2(1, 0);
  for (int i = 1; i <= n; ++i) {
    cout << ans[i] << ' ';
  }
  return 0;
}
2023/11/17 11:46
加载中...