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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <random>
#include <set>
#include <vector>
using namespace std;

#define REP(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define For(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define ALL(v) v.begin(), v.end()

const int NOT_DEFINED = -111;

class BinaryTree {
  int n;
  vector<vector<int>> adjlist;
  vector<vector<int>> inv_adjlist;
  vector<int> depthList;

public:
  BinaryTree(int num);
  int getParent(int id);
  void addNode(int parentId, int childId);
  int getDepth(int id);
  string getType(int id);

  int getDegree(int id);
  int getSibling(int id);
  void fillHeightList(vector<int> &v);
};

BinaryTree::BinaryTree(int num) {
  n = num;
  REP(i, num) {
    vector<int> v;
    adjlist.push_back(v);
    inv_adjlist.push_back(v);
    depthList.push_back(NOT_DEFINED);
  }
}

void BinaryTree::addNode(int parentId, int childId) {
  if (childId == -1)
    return;
  adjlist[parentId].push_back(childId);
  inv_adjlist[childId].push_back(parentId);
}

int BinaryTree::getParent(int id) {
  vector<int> parent = inv_adjlist[id];
  if (parent.size())
    return parent[0];
  return -1;
}

int BinaryTree::getDepth(int id) {
  // if (depthList[id] != NOT_DEFINED) return depthList[id];
  int depth = 0;
  vector<int> parent = inv_adjlist[id];
  while (parent.size()) {
    if (depthList[parent[0]] != NOT_DEFINED) {
      depth += depthList[parent[0]] + 1;
      break;
    }
    parent = inv_adjlist[parent[0]];
    depth++;
  }
  depthList[id] = depth;
  return depth;
}

string BinaryTree::getType(int id) {
  if (depthList[id] == 0) {
    return "root";
  } else if (adjlist[id].empty()) {
    return "leaf";
  } else {
    return "internal node";
  }
  return "";
}

int BinaryTree::getDegree(int id) { return adjlist[id].size(); }

int BinaryTree::getSibling(int id) {
  int parent = getParent(id);
  if (parent == -1)
    return -1;
  if (adjlist[parent].size() == 1)
    return -1;
  REP(i, 2) {
    if (adjlist[parent][i] != id)
      return adjlist[parent][i];
  }
  return -1;
}

void BinaryTree::fillHeightList(vector<int> &v) {
  vector<int> leaves;
  set<int> parents;
  while (1) {
    if (parents.size()) {
      leaves.clear();
      auto itr = parents.begin();
      REP(i, parents.size()) {
        leaves.push_back(*itr);
        itr++;
      }
    } else {
      REP(i, n) {
        if (adjlist[i].size() == 0)
          leaves.push_back(i);
      }
    }

    if (leaves.size() == 1 && inv_adjlist[leaves[0]].size() == 0)
      return;

    parents.clear();
    REP(i, leaves.size()) {
      if (inv_adjlist[leaves[i]].size() > 0) {
        int parent = inv_adjlist[leaves[i]][0];
        parents.insert(parent);
      }
    }
    auto itr = parents.begin();
    REP(i, parents.size()) {
      if (adjlist[*itr].size() == 2)
        v[*itr] = max(v[adjlist[*itr][0]], v[adjlist[*itr][1]]) + 1;
      else
        v[*itr] = v[adjlist[*itr][0]] + 1;
      itr++;
    }
    if (parents.size() == 1 && inv_adjlist[*(parents.begin())].size() == 0)
      return;
  }
}

int main() {
  int n;
  cin >> n;
  BinaryTree tree = BinaryTree(n);
  REP(i, n) {
    int id, deg = 2;
    cin >> id;
    int k;
    REP(j, deg) {
      cin >> k;
      tree.addNode(id, k);
    }
  }

  vector<int> parentList;
  vector<int> siblingList;
  vector<int> degreeList;
  vector<int> depthList;

  REP(i, n) {
    parentList.push_back(tree.getParent(i));
    siblingList.push_back(tree.getSibling(i));
    degreeList.push_back(tree.getDegree(i));
    depthList.push_back(tree.getDepth(i));
  }

  vector<int> heightList(n, 0);
  tree.fillHeightList(heightList);

  REP(i, n) {
    cout << "node " << i << ": parent = ";
    cout << parentList[i] << ", sibling = ";
    cout << siblingList[i] << ", degree = ";
    cout << degreeList[i] << ", depth = ";
    cout << depthList[i] << ", height = ";
    cout << heightList[i] << ", ";
    cout << tree.getType(i) << endl;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <iostream>
#include <random>
#include <set>
#include <vector>
using namespace std;

#define REP(i, n) for (int(i) = 0; (i) < (n); (i)++)
#define For(i, a, b) for (int(i) = (a); (i) < (b); (i)++)
#define ALL(v) v.begin(), v.end()

const int NOT_DEFINED = -111;

class BinaryTree {
  int n;
  vector<vector<int>> adjlist;
  vector<vector<int>> inv_adjlist;
  vector<int> depthList;

public:
  BinaryTree(int num);
  int getParent(int id);
  void addNode(int parentId, int childId);
  int getDepth(int id);
  string getType(int id);

  int getDegree(int id);
  int getSibling(int id);
  void fillHeightList(vector<int> &v);
};

BinaryTree::BinaryTree(int num) {
  n = num;
  REP(i, num) {
    vector<int> v;
    adjlist.push_back(v);
    inv_adjlist.push_back(v);
    depthList.push_back(NOT_DEFINED);
  }
}

void BinaryTree::addNode(int parentId, int childId) {
  if (childId == -1)
    return;
  adjlist[parentId].push_back(childId);
  inv_adjlist[childId].push_back(parentId);
}

int BinaryTree::getParent(int id) {
  vector<int> parent = inv_adjlist[id];
  if (parent.size())
    return parent[0];
  return -1;
}

int BinaryTree::getDepth(int id) {
  // if (depthList[id] != NOT_DEFINED) return depthList[id];
  int depth = 0;
  vector<int> parent = inv_adjlist[id];
  while (parent.size()) {
    if (depthList[parent[0]] != NOT_DEFINED) {
      depth += depthList[parent[0]] + 1;
      break;
    }
    parent = inv_adjlist[parent[0]];
    depth++;
  }
  depthList[id] = depth;
  return depth;
}

string BinaryTree::getType(int id) {
  if (depthList[id] == 0) {
    return "root";
  } else if (adjlist[id].empty()) {
    return "leaf";
  } else {
    return "internal node";
  }
  return "";
}

int BinaryTree::getDegree(int id) { return adjlist[id].size(); }

int BinaryTree::getSibling(int id) {
  int parent = getParent(id);
  if (parent == -1)
    return -1;
  if (adjlist[parent].size() == 1)
    return -1;
  REP(i, 2) {
    if (adjlist[parent][i] != id)
      return adjlist[parent][i];
  }
  return -1;
}

void BinaryTree::fillHeightList(vector<int> &v) {
  vector<int> leaves;
  set<int> parents;
  while (1) {
    if (parents.size()) {
      leaves.clear();
      auto itr = parents.begin();
      REP(i, parents.size()) {
        leaves.push_back(*itr);
        itr++;
      }
    } else {
      REP(i, n) {
        if (adjlist[i].size() == 0)
          leaves.push_back(i);
      }
    }

    parents.clear();
    REP(i, leaves.size()) {
      if (inv_adjlist[leaves[i]].size() > 0) {
        int parent = inv_adjlist[leaves[i]][0];
        parents.insert(parent);
      }
    }
    auto itr = parents.begin();
    REP(i, parents.size()) {
      if (adjlist[*itr].size() == 2)
        v[*itr] = max(v[adjlist[*itr][0]], v[adjlist[*itr][1]]) + 1;
      else
        v[*itr] = v[adjlist[*itr][0]] + 1;
      itr++;
    }
    if (parents.size() == 1)
      return;
  }
}

int main() {
  int n;
  cin >> n;
  BinaryTree tree = BinaryTree(n);
  REP(i, n) {
    int id, deg = 2;
    cin >> id;
    int k;
    REP(j, deg) {
      cin >> k;
      tree.addNode(id, k);
    }
  }

  vector<int> parentList;
  vector<int> siblingList;
  vector<int> degreeList;
  vector<int> depthList;

  REP(i, n) {
    parentList.push_back(tree.getParent(i));
    siblingList.push_back(tree.getSibling(i));
    degreeList.push_back(tree.getDegree(i));
    depthList.push_back(tree.getDepth(i));
  }

  vector<int> heightList(n, 0);
  tree.fillHeightList(heightList);

  REP(i, n) {
    cout << "node " << i << ": parent = ";
    cout << parentList[i] << ", sibling = ";
    cout << siblingList[i] << ", degree = ";
    cout << degreeList[i] << ", depth = ";
    cout << depthList[i] << ", height = ";
    cout << heightList[i] << ", ";
    cout << tree.getType(i) << endl;
  }
  return 0;
}