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

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define irep(i, a, b) for (typeof(i) i = (a); i < (b); ++i)
#define itrep(it, c) for (auto it = (c).begin(); it != (c).end(); ++it)

class Node;

vector<Node> nodes;
ll nn; // チェック済み（戻り済み）ノードの数
map<int, ll>
    excl; // (color, local counter): colorにとって、既に除外されたノード数
map<int, ll> cumm; // (color, paths with no color):
                   // 色が途切れたら、count個数を見てこの領域のpath数を入れる。

class Node {
private:
  int color_;
  set<Node *> links_;

public:
  void setColor(int c) { color_ = c; }
  void setLink(int n) { links_.insert(&nodes[n]); }
  void calc(Node *prev) { // 戻り値はカウント済み全ノード数
    ll mynn = nn;
    ll myexcl = excl[color_];
    itrep(it, links_) {
      if (*it == prev)
        continue;
      (*it)->calc(this); // 戻ってきたら、自分のcolor_については一旦切る。
      ll inner = nn - mynn - (excl[color_] - myexcl);
      cumm[color_] += inner * (inner + 1) / 2;
      excl[color_] += inner;
      //			cout << "INN  color = " << color_ +1 << ", cumm
      //= " << cumm[color_] << ", excl = " << excl[color_] << ", nn = " << nn <<
      //endl;
    }
    nn++;
    excl[color_]++;
    //		cout << "PREV color = " << color_ + 1 << ", cumm = " << cumm[color_]
    //<< ", excl = " << excl[color_] << ", nn = " << nn << endl;
    return;
  }
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);

  ll N;
  cin >> N;
  nodes.assign(N, Node());

  int id = 0;
  itrep(it, nodes) {
    int c;
    cin >> c;
    it->setColor(c - 1);
    id++;
  }

  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    nodes[a - 1].setLink(b - 1);
    nodes[b - 1].setLink(a - 1);
  }

  // loop
  nn = 0;
  nodes.at(0).calc(nullptr);

  rep(i, N) {
    // 最後の加算がされていない
    ll inner = nn - excl[i];
    cumm[i] += inner * (inner + 1) / 2;
    excl[i] = nn;
    //		cout << "color = " << i+1 << ", cumm = " << cumm[i] << ", excl = " <<
    //excl[i] << endl;
    cout << N * (N + 1) / 2 - cumm[i] << endl;
  }
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define irep(i, a, b) for (typeof(i) i = (a); i < (b); ++i)
#define itrep(it, c) for (auto it = (c).begin(); it != (c).end(); ++it)

class Node;
/*
class Tr {
private:
public:
        Tr() {}
        Tr(ll op, ll pa) : opens_(op), paths_(pa) {}
        ll opens_; // 自分のところまで伸びてきているopen pathの数を入れる。
        ll paths_; // 自分を含む、自分より下にあるpathの数を入れる。
        Tr print() {
                cout << "Tr(" << opens_ << "," << paths_ << ")" << endl;
                return(*this);
        }
};
*/

vector<Node> n;

class Node {
private:
  int id_;
  int color_;
  set<int> links_;

public:
  void setId(int id) { id_ = id; }
  void setColor(int c) { color_ = c; }
  void setLink(int l) { links_.insert(l); }
  pair<ll, ll> calc(int color, int prev) {
    int nonCol = (color == color_ ? 0 : 1);
    if (links_.size() == 1 && prev >= 0) {
      //			cout << "1 (opens,closes) = (" << nonCol << ","
      //<< 0 << ")" << endl;
      return (pair<ll, ll>(nonCol, 0));
    } else if (links_.size() == 2 && prev >= 0) {
      ll numOpens = 0;
      ll numCloses = 0;
      ll numOxO = 0;
      pair<ll, ll> t;
      if (*links_.begin() == prev) {
        // 一つ目がprev
        t = n[*links_.rbegin()].calc(color, this->id_);
      } else {
        // 二つ目がprev
        t = n[*links_.begin()].calc(color, this->id_);
      }
      pair<ll, ll> ret =
          pair<ll, ll>(nonCol * (t.first + 1), t.first + t.second);
      //			cout << "2 (opens,closes) = (" << ret.first <<
      //"," << ret.second << ")" << endl;
      return (ret);
    } else {
      ll numOpens = 0;
      ll numCloses = 0;
      ll numOxO = 0;
      itrep(it, links_) {
        if (*it == prev) {
          continue;
        }
        pair<ll, ll> t = n[*it].calc(color, this->id_);
        numOpens += t.first;
        numCloses += t.second;
        numOxO += t.first * t.first;
      }
      pair<ll, ll> ret = pair<ll, ll>(
          nonCol * (numOpens + 1),
          numOpens + numCloses + nonCol * (numOpens * numOpens - numOxO) / 2);
      //			cout << "(opens,closes) = (" << ret.first << ","
      //<< ret.second << ")" << endl;
      return (ret);
    }
  }
};

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);

  ll N;
  cin >> N;
  n.assign(N, Node());
  set<int> cs; // 色の種類を入れる。

  int id = 0;
  itrep(it, n) {
    int c;
    cin >> c;
    it->setId(id);
    it->setColor(c);
    cs.insert(c);
    id++;
  }

  rep(i, N - 1) {
    int a, b;
    cin >> a >> b;
    n[a - 1].setLink(b - 1);
    n[b - 1].setLink(a - 1);
  }

  // loop
  rep(i, N) {
    if (cs.find(i + 1) != cs.end()) {
      pair<ll, ll> ret = n[0].calc(i + 1, -1);
      //			cout << "color = " << i << ", paths = " <<
      //ret.paths_ << endl;
      cout << N * (N + 1) / 2 - ret.second - ret.first << endl;
    } else {
      //			cout << "color = " << i << ", paths = " << 0 <<
      //endl;
      cout << 0 << endl;
    }
  }
}