Problem p01884 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <set>
#include <string>
using namespace std;
struct node {
  int x, y;
} a[200010];
int ednum[100010];
int flag[100010];
unsigned long long dis[100010];
bool cmp(const node &a, const node &b) {
  if (a.x == b.x)
    return a.y < b.y;
  return a.x < b.x;
}
unsigned long long dfs(long long x) {
  long long j;
  unsigned long long ans = 0, num = 0;
  for (j = ednum[x - 1]; j < ednum[x]; j++)
    if (flag[a[j].y] == 0) {
      flag[a[j].y] = 1;
      num = num + 1;
      ans = ans + dfs(a[j].y);
      flag[a[j].y] = 0;
    }
  dis[x] = num + ans * 100003;
  return dis[x];
}
int main() {
  long long i, j, n, len, x, num, ans;
  while (~scanf("%lld", &n)) {
    memset(ednum, 0, sizeof(ednum));
    for (i = 0; i < n - 1; i++) {
      scanf("%d%d", &a[i].x, &a[i].y);
      a[i + n - 1].x = a[i].y;
      a[i + n - 1].y = a[i].x;
      ednum[a[i].x]++;
      ednum[a[i].y]++;
    }
    for (i = 1; i <= n; i++)
      ednum[i] += ednum[i - 1];
    sort(a, a + 2 * n - 2, cmp);
    memset(dis, 0, sizeof(dis));
    memset(flag, 0, sizeof(flag));
    flag[1] = 1;
    dfs(1);
    sort(dis, dis + n);
    ans = 0;
    for (i = 1, num = 1; i < n; i++) {
      if (dis[i] == dis[i - 1])
        num++;
      else {
        ans = ans + (num - 1) * num / 2;
        num = 1;
      }
    }
    ans = ans + (num - 1) * num / 2;
    /*		for(i=1;i<=n;i++)
                    {
                            printf("i:%lld dis:%llu\\n",i,dis[i]);
                    }*/
    printf("%lld\\n", ans);
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <queue>
#include <set>
#include <vector>
#pragma comment(linker, "/STACK:1024000000,1024000000")
using namespace std;
vector<int> tu[100100];
queue<int> have;
int vis[100100];
int dp[100100];
set<int> maxdep;
int fu[100100];
int nedhai[100100];
int nowhai[100100];
class YUAN {
public:
  int len;
  vector<int> node;

public:
  void init() {
    len = 0;
    node.clear();
  }
  bool operator<(const YUAN &t) const {
    if (len != t.len)
      return len < t.len;
    for (int i = 0; i < len; i++) {
      if (node[i] != t.node[i])
        return node[i] < t.node[i];
    }
    return len < t.len;
  }
  bool operator==(const YUAN &t) const {
    if (len != t.len)
      return len == t.len;
    for (int i = 0; i < len; i++) {
      if (node[i] != t.node[i])
        return node[i] == t.node[i];
    }
    return len == t.len;
  }
} yuan[100100], all[100100];
long long c(int x) {
  return (long long)x * (long long)(x - 1ll) / (long long)2;
}
void dfs(int root) {
  // printf("%d\\n",root);
  if (yuan[root].len == 0) {
    yuan[root].node.push_back(1);
    yuan[root].len++;
  }
  vis[root] = 1;
  int siz = tu[root].size();
  for (int i = 0; i < siz; i++) {
    int new_node = tu[root][i];
    if (vis[new_node] != 0)
      continue;
    dfs(new_node);
    vis[root] = max(vis[root], vis[new_node] + 1);
  }
  if (vis[root] == 1) {
    maxdep.insert(root);
  }
  return;
}
long long solve() {
  long long res = 0ll;
  while (1) {
    while (have.size() != 0)
      have.pop();
    set<int>::iterator it;
    if (maxdep.size() == 0)
      break;
    int cnt = 1;
    for (it = maxdep.begin(); it != maxdep.end(); it++) {
      int now = *it;
      if (nowhai[now] != nedhai[now])
        continue;
      all[cnt++] = yuan[now];
    }
    sort(all + 1, all + cnt);
    dp[1] = 1;
    for (int i = 2; i <= cnt - 1; i++) {
      if (all[i] == all[i - 1]) {
        dp[i] = dp[i - 1] + 1;
      } else {
        dp[i] = 1;
        res += c(dp[i - 1]);
      }
    }
    res += c(dp[cnt - 1]);
    for (it = maxdep.begin(); it != maxdep.end(); it++) {
      int now = *it;
      int fuqin = fu[now];
      if (fuqin == 0) {
        yuan[now].init();
        have.push(now);
        continue;
      }
      if (nowhai[now] != nedhai[now])
        continue;
      while (yuan[fuqin].len <= yuan[now].len) {
        yuan[fuqin].node.push_back(0);
        yuan[fuqin].len++;
      }
      for (int i = 0; i < yuan[now].len; i++) {
        yuan[fuqin].node[i + 1] += yuan[now].node[i];
      }
      nowhai[fuqin]++;
      yuan[now].init();
      have.push(now);
      maxdep.insert(fuqin);
    }
    while (have.size() != 0) {
      int x = have.front();
      maxdep.erase(x);
      have.pop();
    }
  }
  return res;
}
int main() {
  // freopen("E://input.txt","r",stdin);
  int n;
  while (scanf("%d", &n) != EOF) {
    for (int i = 1; i <= n + 10; i++) {
      tu[i].clear();
      yuan[i].init();
      fu[i] = 0;
      nowhai[i] = 0;
      nedhai[i] = 0;
    }
    for (int i = 1; i < n; i++) {
      int x, y;
      scanf("%d%d", &x, &y);
      tu[x].push_back(y);
      tu[y].push_back(x);
      fu[y] = x;
      nedhai[x]++;
    }
    memset(vis, 0, sizeof(vis));
    dfs(1);
    long long res = solve();
    printf("%lld\\n", res);
  }
}