Skip to content

Instantly share code, notes, and snippets.

@NamPE286
Last active April 16, 2023 12:29
Show Gist options
  • Save NamPE286/f93ba90725f7ba5075c71c22df7e3857 to your computer and use it in GitHub Desktop.
Save NamPE286/f93ba90725f7ba5075c71c22df7e3857 to your computer and use it in GitHub Desktop.
Dynamic general tree class implementation in C++
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
template <typename T>
class general_tree {
private:
T root;
bool modified = false;
unordered_map<T, vector<T>> graph;
unordered_map<T, T> parent;
unordered_map<T, ll> subtree_size, depth;
void pre_calc_size(ll cur_node) {
subtree_size[cur_node] = 1;
for (T i : graph[cur_node]) {
if (i == parent[cur_node]) continue;
depth[cur_node]++;
pre_calc_size(i);
subtree_size[cur_node] += subtree_size[i];
}
}
public:
general_tree(T tree_root) {
root = tree_root;
}
void connect(T parent_node, T children_node) {
parent[children_node] = parent_node;
graph[parent_node].push_back(children_node);
graph[children_node].push_back(parent_node);
modified = true;
}
ll size(T node) {
if (modified) {
pre_calc_size(root);
subtree_size.clear();
depth.clear();
modified = false;
}
return subtree_size[node];
}
};
int main() {
ll n;
cin >> n;
general_tree<ll> tree(1);
for (ll i = 2; i <= n; i++) {
ll x;
cin >> x;
tree.connect(x, i);
}
for (ll i = 1; i <= n; i++) {
cout << tree.size(i) - 1 << ' ';
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment