Submission #1143678

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143678		JelalTkm	Balanced Tree (info1cup18_balancedtree)	C++20	0 / 100	51 ms	7204 KiB

balancedtree

#include <bits/stdc++.h>
#pragma GCC optimize ("O3")
#pragma GCC target ("sse4")

using namespace std;

#define int long long int

const int N = 1e3 + 1;
const int md = 1e9 + 7;
const int INF = 1e9;

int32_t main(int32_t argc, char *argv[]) {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

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

    vector<int> a(n + 1);
    for (int i = 1; i <= n; i++)
      cin >> a[i];

    vector<vector<int>> dp(n + 1, vector<int> (2, INF)), as(n + 1, vector<int> (2, INF)),
    dp1(n + 1, vector<int> (2, INF)), av(n + 1, vector<int> (2, INF));
    function<void(int, int)> dfs = [&](int u, int p) {
      bool ok = 1;
      dp1[u][a[u]] = 0;
      for (auto v : g[u])
        if (v != p) {
          dp1[v][0] = dp1[u][0] + 1;
          dp1[v][1] = dp1[u][1] + 1;
          av[v][0] = dp1[u][0] + 1;
          av[v][1] = dp1[u][1] + 1;
          dfs(v, u);
          dp[u][0] = dp[v][0] + 1;
          dp[u][1] = dp[v][1] + 1;
          as[u][0] = dp[v][0] + 1;
          as[u][1] = dp[v][1] + 1;
          ok = 0;
        }
      dp[u][a[u]] = 0;
      return;
    };

    dfs(1, -1);

    vector<vector<int>> ans(n + 1, vector<int> (2, INF));
    function<void(int, int)> dfs1 = [&](int u, int p) {
      bool ok = 1;
      vector<vector<int>> vi;
      for (auto v : g[u])
        if (v != p) {
          vi.push_back({min(av[v][0], as[v][0]), min(av[v][1], as[v][1])});
        }
      int sz = (int) vi.size();
      if (sz) {
        vector<vector<int>> pref(sz, vector<int> {INF, INF}), suff(sz + 1, vector<int> {INF, INF});
        pref[0] = vi[0];
        for (int i = 1; i < sz; i++) {
          pref[i] = pref[i - 1];
          pref[i][0] = min(pref[i][0], vi[i][0]);
          pref[i][1] = min(pref[i][1], vi[i][1]);
        }
        suff[0] = vi[sz - 1];
        for (int i = sz - 2; i >= 0; i--) {
          suff[i] = suff[i + 1];
          suff[i][0] = min(suff[i][0], vi[i][0]);
          suff[i][1] = min(suff[i][1], vi[i][1]);
        }
        int cnt = 0;
        for (auto v : g[u])
          if (v != p) {
            ans[v][0] = min((cnt > 0 ? pref[cnt - 1][0] : INF), suff[cnt + 1][0]) + 2;
            ans[v][1] = min((cnt > 0 ? pref[cnt - 1][1] : INF), suff[cnt + 1][1]) + 2;
            cnt++;
          }
      }
      for (auto v : g[u])
        if (v != p) {
          ans[v][0] = min(ans[v][0], av[v][0]);
          ans[v][1] = min(ans[v][1], av[v][1]);
          dfs1(v, u);
        }
      return;
    };

    dfs1(1, -1);

    ans[1][a[1]] = as[1][a[1]];

    int mx = 0;
    for (int i = 1; i <= n; i++)
      mx = max(mx, min({av[i][a[i]], as[i][a[i]]}));

    if (mx >= INF) {
      cout << -1 << '\n';
      continue;
    }

    cout << mx << '\n';
    for (int i = 1; i <= n; i++)
      cout << a[i] << " ";
    cout << '\n';
  }

  return 0;
}

• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 26
• Subtask 4: 0 / 45
• Subtask 5: 0 / 6