树模板

struct Tree {
   int n, d = 0; //顶点,直径,树中心
   //自顶向下u到其叶子结点最远距离d1,次长距离d2(与最长路径无公共边)
   //p1,p2表示结点u向下更新时是由哪个结点更新来的
   //up表示结点u向上到祖宗结点的最远距离
   vector<int> d1, d2, p1, p2, up;
   vector<vector<int>> g;
   Tree(int n): n(n), d1(n + 1), d2(n + 1), g(n + 1), p1(n + 1), p2(n + 1), up(n + 1) {}

   void  add(int u, int v) {
      g[u].emplace_back(v);
      g[v].emplace_back(u);
   }
   //求直径//自顶向下求u到叶子结点的最远距离
   void dfs(int u, int fa) {
      d1[u] = d2[u] = 0;
      for (auto v : g[u]) {
         if (v == fa)continue;
         dfs(v, u);
         auto t = d1[v] + 1;
         if (t > d1[u]) {
            d2[u] = d1[u], p2[u] = p1[u];
            d1[u] = t, p1[u] = v;
         } else if (t > d2[u]) {
            d2[u] = t, p2[u] = v;
         }
      }
      d = max(d, d1[u] + d2[u]);
   }
   //自底向上求u到其它结点的最长路径
   void dfsup(int u, int fa) {
      for (auto v : g[u]) {
         if (v == fa) continue;
         //如果父结点u向下的最长路径经过v
         if (p1[u] == v) {
            //结点v向上走到最长路径为
            //父结点u继续向上的的最长路径和u向下走的次长路径的最大值+边权
            up[v] = max(up[u], d2[u]) + 1;
         } else {
            //如果父结点u向下的最长路径不经过v
            up[v] = max(up[u], d1[u]) + 1;
         }
         dfsup(v, u);
      }
   }

   //求树的直径
   int FindDiameter() {
      dfs(1, 0);
      return d;
   }

   //求树的中心
   vector<int> FindCenter() {
      dfsup(1, 0);
      i64 res = INT_MAX;
      vector<int> mid;
      for (int i = 1; i <= n; i ++) {
         auto t = max(d1[i], up[i]);
         if (t < res) {
            res = t;
            vector<int>().swap(mid);
            mid.emplace_back(i);
         } else if (t == res) {
            mid.emplace_back(i);
         }
      }
      return mid;
   }

};
posted @ 2024-07-27 16:41  Ke_scholar  阅读(7)  评论(0编辑  收藏  举报