Problem p02710 (C++) — buggy submission, verdict: Time Limit Exceeded


/*
Writer: SPD_9X2
https://atcoder.jp/contests/abc163/custom_test
どう考えても余事象
つまり、各色に関してそれを通らないパスの数を計算できれば良い

1色のみを考える場合→いい感じにdfs
自分がその色なら、各子について nC2をansに加えて自分は0を返す
問題は色が多い事

オイラーツアーと聞いた
オイラーツアーをすると尺取り法で求められる？
葉でreturnするときに、2つappendすることで、おなじ頂点が2個
入るのを防ぐことが出来る

=====解法を確認======

公式開放？？？？
なのでマージテクの方でいこう

ans[c] = 色cを通らないパスの数
chnum[v] = vの部分木にある頂点数
dp[v][c] = vの部分木の内,根から色cを通らないと行けない頂点数

とすれば、各v(色c)について、各子uに関して
ans[c] += nCr( chnum[u] - dp[u][c] , 2 )
とし、dp[v][c] = 0とすればよい

dp[v][c]のcの部分をmapにすれば、マージテクで、マージはO(NlogN)になる。
探索がO(logN)なので,結果 O(NlogNlogN)で解ける

*/

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n, m) for (ll(i) = (n); (i) < (m); (i)++)
#define rrep(i, n, m) for (ll(i) = (n); (i) > (m); (i)--)
using ll = long long;
const ll mod = 998244353;

ll nC2(ll x) { return x * (x - 1) / 2; }

void dfs(ll v, ll p, vector<vector<ll>> &lis, vector<ll> &col, vector<ll> &ans,
         vector<ll> &chnum, vector<map<ll, ll>> &dp) {

  rep(i, 0, lis[v].size()) {
    if (lis[v][i] != p) {
      dfs(lis[v][i], v, lis, col, ans, chnum, dp);
      ans[col[v]] += nC2(chnum[lis[v][i]] - dp[lis[v][i]][col[v]]);
      chnum[v] += chnum[lis[v][i]];

      if (dp[v].size() < dp[lis[v][i]].size()) {
        swap(dp[v], dp[lis[v][i]]);
      }

      for (const auto &[k, val] : dp[lis[v][i]]) {
        dp[v][k] = dp[v][k] + val;
      }
    }
  }

  dp[v][col[v]] = chnum[v];

  return;
}

int main() {

  ll N;
  cin >> N;
  vector<ll> c(N, 0);
  rep(i, 0, N) {
    cin >> c[i];
    c[i]--;
  }
  vector<vector<ll>> lis(N, vector<ll>(0));
  ll a, b;
  rep(i, 0, N - 1) {
    cin >> a >> b;
    a--;
    b--;
    lis[a].push_back(b);
    lis[b].push_back(a);
  }

  vector<ll> ans(N, 0);
  vector<ll> chnum(N, 1);
  vector<map<ll, ll>> dp(0);
  rep(i, 0, N) {
    map<ll, ll> tmp;
    dp.push_back(tmp);
  }

  dfs(0, 0, lis, c, ans, chnum, dp);
  rep(i, 0, N) { ans[i] += nC2(N - dp[0][i]); }

  rep(i, 0, N) { ans[i] = nC2(N) - ans[i]; }

  rep(i, 0, N) { ans[c[i]]++; }

  rep(i, 0, N) { cout << ans[i] << endl; }
}

Error: TLE

Fixed submission:


/*
Writer: SPD_9X2
https://atcoder.jp/contests/abc163/custom_test
どう考えても余事象
つまり、各色に関してそれを通らないパスの数を計算できれば良い

1色のみを考える場合→いい感じにdfs
自分がその色なら、各子について nC2をansに加えて自分は0を返す
問題は色が多い事

オイラーツアーと聞いた
オイラーツアーをすると尺取り法で求められる？
葉でreturnするときに、2つappendすることで、おなじ頂点が2個
入るのを防ぐことが出来る

=====解法を確認======

公式開放？？？？
なのでマージテクの方でいこう

ans[c] = 色cを通らないパスの数
chnum[v] = vの部分木にある頂点数
dp[v][c] = vの部分木の内,根から色cを通らないと行けない頂点数

とすれば、各v(色c)について、各子uに関して
ans[c] += nCr( chnum[u] - dp[u][c] , 2 )
とし、dp[v][c] = 0とすればよい

dp[v][c]のcの部分をmapにすれば、マージテクで、マージはO(NlogN)になる。
探索がO(logN)なので,結果 O(NlogNlogN)で解ける

*/

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n, m) for (ll(i) = (n); (i) < (m); (i)++)
#define rrep(i, n, m) for (ll(i) = (n); (i) > (m); (i)--)
using ll = long long;
const ll mod = 998244353;

ll nC2(ll x) { return x * (x - 1) / 2; }

map<ll, ll> mgmap(map<ll, ll> a, map<ll, ll> b) {
  if (a.size() > b.size())
    swap(a, b);
  for (const auto &[k, v] : a) {
    b[k] = b[k] + v;
  }
  return b;
}

void dfs(ll v, ll p, vector<vector<ll>> &lis, vector<ll> &col, vector<ll> &ans,
         vector<ll> &chnum, vector<map<ll, ll>> &dp) {

  rep(i, 0, lis[v].size()) {
    if (lis[v][i] != p) {
      dfs(lis[v][i], v, lis, col, ans, chnum, dp);
      dp[v] = mgmap(dp[v], dp[lis[v][i]]);
      chnum[v] += chnum[lis[v][i]];

      ans[col[v]] += nC2(chnum[lis[v][i]] - dp[lis[v][i]][col[v]]);
    }
  }

  dp[v][col[v]] = chnum[v];

  return;
}

int main() {

  ll N;
  cin >> N;
  vector<ll> c(N, 0);
  rep(i, 0, N) {
    cin >> c[i];
    c[i]--;
  }
  vector<vector<ll>> lis(N, vector<ll>(0));
  ll a, b;
  rep(i, 0, N - 1) {
    cin >> a >> b;
    a--;
    b--;
    lis[a].push_back(b);
    lis[b].push_back(a);
  }

  vector<ll> ans(N, 0);
  vector<ll> chnum(N, 1);
  vector<map<ll, ll>> dp(0);
  rep(i, 0, N) {
    map<ll, ll> tmp;
    dp.push_back(tmp);
  }

  dfs(0, 0, lis, c, ans, chnum, dp);
  rep(i, 0, N) { ans[i] += nC2(N - dp[0][i]); }

  rep(i, 0, N) { ans[i] = nC2(N) - ans[i]; }

  rep(i, 0, N) { ans[c[i]]++; }

  rep(i, 0, N) { cout << ans[i] << endl; }
}