some programs

JLOI2009_binary_tree.cpp

@@ -0,0 +1,87 @@
+#include <iostream>
+#include <vector>
+#include <cassert>
+
+using namespace std;
+
+class Graph {
+public:
+    struct Node {
+        vector<int> neighbors;
+        int father = -1;
+        int depth = -1;
+    };
+
+    vector<Node> nodes;
+    int depth = -1, breadth = -1;
+
+    void connect(int x, int y) {
+        assert(x < nodes.size() && y < nodes.size());
+        nodes[x].neighbors.push_back(y);
+        nodes[y].neighbors.push_back(x);
+    }
+
+    void init_tree() {
+        vector<int> depths_map;
+        init_tree(0, 0, depths_map);
+        depth = depths_map.size();
+        breadth = -1;
+        for (const auto& d : depths_map) {
+            breadth = max(breadth, d);
+        }
+    }
+
+    void init_tree(int node_idx, int current_depth, vector<int>& depths_map) {
+        if (depths_map.size() <= current_depth) { depths_map.push_back(0); }
+        depths_map[current_depth]++;
+        for (const auto& neighbor : nodes[node_idx].neighbors) {
+            if (neighbor == nodes[node_idx].father) { continue; }
+            nodes[neighbor].father = node_idx;
+            nodes[neighbor].depth = current_depth + 1;
+            init_tree(neighbor, current_depth + 1, depths_map);
+        }
+    }
+
+    int get_cost(int start, int end) {
+        int cost = 0;
+        while (nodes[start].depth > nodes[end].depth) {
+            cost += 2;
+            start = nodes[start].father;
+        }
+        while (nodes[start].depth < nodes[end].depth) {
+            cost += 1;
+            end = nodes[end].father;
+        }
+        while (start != end) {
+            cost += 3;
+            start = nodes[start].father;
+            end = nodes[end].father;
+        }
+        return cost;
+    }
+};
+
+
+Graph tree;
+int n;
+
+int main() {
+    ios::sync_with_stdio(false);
+    cin.tie(nullptr);
+
+    int x, y;
+
+    cin >> n;
+    tree.nodes.resize(n);
+    for (int _ = 1; _ < n; _++) {
+        cin >> x >> y;
+        tree.connect(x - 1, y - 1);
+    }
+    tree.init_tree();
+    cout << tree.depth << endl << tree.breadth << endl;
+
+    cin >> x >> y;
+    cout << tree.get_cost(x - 1, y - 1) << endl;
+
+    return 0;
+}