#1000092 (C++23) No.898 tri-βutree

提出ソース
結果

問題	No.898 tri-βutree
ユーザー	KowerKoint2010
提出日時	2024-08-06 17:36:06
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	551 ms / 4,000 ms
コード長	19,331 bytes
コンパイル時間	3,779 ms
コンパイル使用メモリ	275,108 KB
実行使用メモリ	41,240 KB
最終ジャッジ日時	2024-08-06 17:36:22
合計ジャッジ時間	15,420 ms
ジャッジサーバーID （参考情報）	judge4 / judge5
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	312 ms 41,240 KB
testcase_01	AC	3 ms 5,376 KB
testcase_02	AC	3 ms 5,376 KB
testcase_03	AC	3 ms 5,376 KB
testcase_04	AC	3 ms 5,376 KB
testcase_05	AC	2 ms 5,376 KB
testcase_06	AC	3 ms 5,376 KB
testcase_07	AC	551 ms 40,172 KB
testcase_08	AC	508 ms 40,180 KB
testcase_09	AC	525 ms 40,176 KB
testcase_10	AC	518 ms 40,180 KB
testcase_11	AC	533 ms 40,164 KB
testcase_12	AC	516 ms 40,288 KB
testcase_13	AC	536 ms 40,300 KB
testcase_14	AC	519 ms 40,180 KB
testcase_15	AC	536 ms 40,176 KB
testcase_16	AC	508 ms 40,168 KB
testcase_17	AC	506 ms 40,176 KB
testcase_18	AC	520 ms 40,296 KB
testcase_19	AC	507 ms 40,300 KB
testcase_20	AC	541 ms 40,292 KB
testcase_21	AC	511 ms 40,296 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#line 1 "test/yukicoder-898.test.cpp"
#define PROBLEM "https://yukicoder.me/problems/no/898"

#line 2 "cpp/tree.hpp"

/**
 * @file tree.hpp
 * @brief 木の汎用テンプレート
*/

#include <algorithm>
#include <stack>
#include <tuple>

#line 2 "cpp/graph.hpp"

#include <iostream>
#include <limits>
#include <queue>
#include <vector>

/**
 * @brief グラフの汎用クラス
 *
 * @tparam Cost 辺のコストの型
 */
template <typename Cost=int>
struct Graph {
    /**
     * @brief 有向辺の構造体
     *
     * operator int()を定義しているので、int型にキャストすると勝手にdstになる
     * 例えば、
     * for (auto& e : g[v]) をすると、vから出る辺が列挙されるが、
     * for (int dst : g[v]) とすると、vから出る辺の行き先が列挙される
     */
    struct Edge {
        int src; //!< 始点
        int dst; //!< 終点
        Cost cost; //!< コスト
        int id; //!< 辺の番号(追加された順、無向辺の場合はidが同じで方向が逆のものが2つ存在する)
        Edge() = default;
        Edge(int src, int dst, Cost cost=1, int id=-1) : src(src), dst(dst), cost(cost), id(id) {}
        operator int() const { return dst; }
    };

    int n; //!< 頂点数
    int m; //!< 辺数
    std::vector<std::vector<Edge>> g; //!< グラフの隣接リスト表現

    /**
     * @brief デフォルトコンストラクタ
     */
    Graph() : n(0), m(0), g(0) {}
    /**
     * @brief コンストラクタ
     * @param n 頂点数
     */
    explicit Graph(int n) : n(n), m(0), g(n) {}
    /**
     * @brief 無向辺を追加する
     * @param u 始点
     * @param v 終点
     * @param w コスト 省略したら1
     */
    void add_edge(int u, int v, Cost w=1) {
        g[u].push_back({u, v, w, m});
        g[v].push_back({v, u, w, m++});
    }
    /**
     * @brief 有向辺を追加する
     * @param u 始点
     * @param v 終点
     * @param w コスト 省略したら1
     */
    void add_directed_edge(int u, int v, Cost w=1) {
        g[u].push_back({u, v, w, m++});
    }
    /**
     * @brief 辺の情報を標準入力から受け取って追加する
     * @param m 辺の数
     * @param padding 頂点番号を入力からいくつずらすか 省略したら-1
     * @param weighted 辺の重みが入力されるか 省略したらfalseとなり、重み1で辺が追加される
     * @param directed 有向グラフかどうか 省略したらfalse
     */
    void read(int m, int padding=-1, bool weighted=false, bool directed=false) {
        for(int i = 0; i < m; i++) {
            int u, v; std::cin >> u >> v; u += padding, v += padding;
            Cost c(1);
            if(weighted) std::cin >> c;
            if(directed) add_directed_edge(u, v, c);
            else add_edge(u, v, c);
        }
    }
    /**
     * @brief ある頂点から出る辺を列挙する
     * @param v 頂点番号
     * @return std::vector<Edge>& vから出る辺のリスト
     */
    std::vector<Edge>& operator[](int v) {
        return g[v];
    }
    /**
     * @brief ある頂点から出る辺を列挙する
     * @param v 頂点番号
     * @return const std::vector<Edge>& vから出る辺のリスト
     */
    const std::vector<Edge>& operator[](int v) const {
        return g[v];
    }
    /**
     * @brief 辺のリスト
     * @return std::vector<Edge> 辺のリスト(idの昇順)
     *
     * 無向辺は代表して1つだけ格納される
     */
    std::vector<Edge> edges() const {
        std::vector<Edge> res(m);
        for(int i = 0; i < n; i++) {
            for(auto& e : g[i]) {
                res[e.id] = e;
            }
        }
        return res;
    }
    /**
     * @brief ある頂点から各頂点への最短路
     *
     * @param s 始点
     * @param weighted 1以外のコストの辺が存在するか 省略するとtrue
     * @param inf コストのminの単位元 未到達の頂点への距離はinfになる 省略すると-1
     * @return std::pair<std::vector<Cost>, std::vector<Edge>> first:各頂点への最短路長 second:各頂点への最短路上の直前の辺
     */
    std::pair<std::vector<Cost>, std::vector<Edge>> shortest_path(int s, bool weignted = true, Cost inf = -1) const {
        if(weignted) return shortest_path_dijkstra(s, inf);
        return shortest_path_bfs(s, inf);
    }
    
    std::vector<int> topological_sort() {
        std::vector<int> indeg(n), sorted;
        std::queue<int> q;
        for (int i = 0; i < n; i++) {
            for (int dst : g[i]) indeg[dst]++;
        }
        for (int i = 0; i < n; i++) {
            if (!indeg[i]) q.push(i);
        }
        while (!q.empty()) {
            int cur = q.front(); q.pop();
            for (int dst : g[cur]) {
                if (!--indeg[dst]) q.push(dst);
            }
            sorted.push_back(cur);
        }
        return sorted;
    }

private:
    std::pair<std::vector<Cost>, std::vector<Edge>> shortest_path_bfs(int s, Cost inf) const {
        std::vector<Cost> dist(n, inf);
        std::vector<Edge> prev(n);
        std::queue<int> que;
        dist[s] = 0;
        que.push(s);
        while(!que.empty()) {
            int u = que.front(); que.pop();
            for(auto& e : g[u]) {
                if(dist[e.dst] == inf) {
                    dist[e.dst] = dist[e.src] + 1;
                    prev[e.dst] = e;
                    que.push(e.dst);
                }
            }
        }
        return {dist, prev};
    }
    std::pair<std::vector<Cost>, std::vector<Edge>> shortest_path_dijkstra(int s, Cost inf) const {
        std::vector<Cost> dist(n, inf);
        std::vector<Edge> prev(n);
        using Node = std::pair<Cost, int>;
        std::priority_queue<Node, std::vector<Node>, std::greater<Node>> que;
        dist[s] = 0;
        que.push({0, s});
        while(!que.empty()) {
            auto [d, u] = que.top(); que.pop();
            if(d > dist[u]) continue;
            for(auto& e : g[u]) {
                if(dist[e.dst] == inf || dist[e.dst] > dist[e.src] + e.cost) {
                    dist[e.dst] = dist[e.src] + e.cost;
                    prev[e.dst] = e;
                    que.push({dist[e.dst], e.dst});
                }
            }
        }
        return {dist, prev};
    }


};
#line 13 "cpp/tree.hpp"

template <typename>
struct RootedTree;
template <typename>
struct DoublingClimbTree;

/**
 * @brief 根なし木
 * 
 * @tparam Cost = int 辺の重み
 * 
 * Graph<Cost>を継承し、すべて無向辺で表す
*/
template <typename Cost = int>
struct Tree : private Graph<Cost> {
    // Graph<Cost>* g = new Tree<Cost>(); ができてしまうと、delete時にメモリリークが発生
    // 回避するためprivate継承にして、メンバをすべてusing宣言
    using Edge = typename Graph<Cost>::Edge;
    using Graph<Cost>::n;
    using Graph<Cost>::m;
    using Graph<Cost>::g;
    using Graph<Cost>::Graph;
    using Graph<Cost>::add_edge;
    using Graph<Cost>::add_directed_edge;
    using Graph<Cost>::read;
    using Graph<Cost>::operator[];
    using Graph<Cost>::edges;
    using Graph<Cost>::shortest_path;

    /**
     * @brief コンストラクタ
     * 
     * @param n ノード数
     */
    Tree(int n = 0) : Graph<Cost>(n) {}

    /**
     * @brief 辺の情報を標準入力から受け取って追加する
     * @param padding  = -1 頂点番号を入力からいくつずらすか
     * @param weighted = false 辺の重みが入力されるか
     */
    void read(int padding = -1, bool weighted = false) {
        Graph<Cost>::read(this->n - 1, padding, weighted, false);
    }
    
    /**
     * @brief 木の直径
     * 
     * @param weighted = true 重み付きか
     * @return std::vector<Edge> 直径を構成する辺のリスト
     */
    std::vector<Edge> diameter(bool weighted = true) const {
        std::vector<Cost> dist = shortest_path(0, weighted).first;
        int u = std::max_element(dist.begin(), dist.end()) - dist.begin();
        std::vector<Edge> prev;
        std::tie(dist, prev) = shortest_path(u, weighted);
        int v = std::max_element(dist.begin(), dist.end()) - dist.begin();
        std::vector<Edge> path;
        while(v != u) {
            path.push_back(prev[v]);
            v = prev[v].src;
        }
        reverse(path.begin(), path.end());
        return path;
    }

    /**
     * @brief 根付き木にする
     * 
     * @param root 根
     * @return RootedTree<Cost> 根付き木のオブジェクト
     */
    [[nodiscard]]
    RootedTree<Cost> set_root(int root) const& {
        return RootedTree<Cost>(*this, root);
    }
    /**
     * @brief 根付き木にして返す
     * 
     * @param root 根
     * @return RootedTree<Cost> 根付き木のオブジェクト
     */
    [[nodiscard]]
    RootedTree<Cost> set_root(int root) && {
        return RootedTree<Cost>(std::move(*this), root);
    }

    /**
     * @brief ダブリングLCAが使える根付き木を返す
     * 
     * @param root 根
     * @return DoublingClimbTree<Cost> ダブリング済み根付き木のオブジェクト
     */
    [[nodiscard]]
    DoublingClimbTree<Cost> build_lca(int root) const& {
        RootedTree rooted_tree(*this, root);
        return DoublingClimbTree<Cost>(std::move(rooted_tree));
    }
    /**
     * @brief ダブリングLCAが使える根付き木を返す
     * 
     * @param root 根
     * @return DoublingClimbTree<Cost> ダブリング済み根付き木のオブジェクト
     */
    [[nodiscard]]
    DoublingClimbTree<Cost> build_lca(int root) && {
        RootedTree rooted_tree(std::move(*this), root);        
        return DoublingClimbTree<Cost>(std::move(rooted_tree));
    }
};

/**
 * @brief 根付き木
 * 
 * @tparam Cost = int 辺のコスト
 */
template <typename Cost = int>
struct RootedTree : private Tree<Cost> {
    using Edge = typename Tree<Cost>::Edge;
    using Tree<Cost>::n;
    using Tree<Cost>::m;
    using Tree<Cost>::g;
    using Tree<Cost>::operator[];
    using Tree<Cost>::edges;
    using Tree<Cost>::shortest_path;

    int root; //!< 根
    std::vector<Edge> par; //!< 親へ向かう辺
    std::vector<std::vector<Edge>> child; //!< 子へ向かう辺のリスト
    std::vector<Cost> depth; //!< 深さのリスト
    std::vector<Cost> height; //!< 部分木の高さのリスト
    std::vector<int> unweighted_depth; //!< 重みなしの深さのリスト
    std::vector<int> unweighted_height; //!< 重みなしの部分木の高さのリスト
    std::vector<int> sz; //!< 部分期の頂点数のリスト
    std::vector<int> preorder; //!< 先行順巡回
    std::vector<int> postorder; //!< 後行順巡回
    std::vector<std::pair<int, int>> euler_tour; //!< オイラーツアー DFS中にとおった頂点の(頂点番号,その頂点を通った回数)のリスト

    RootedTree() = delete;
    
    /**
     * @brief 親の頂点のリストから根が0の根付き木を構築するコンストラクタ
     * 
     * @param par_ 頂点0以外の親の頂点のリスト
     * @param padding = -1 parの頂点番号をいくつずらすか
     */
    RootedTree(const std::vector<int>& par_, int padding = -1) : Tree<Cost>(par_.size() + 1), root(0) {
        for(int i = 0; i < (int)par_.size(); i++) {
            this->add_edge(i + 1, par_[i] + padding);
        }
        build();
    }
    /**
     * @brief Treeから根付き木を構築するコンストラクタ(コピー)
     *
     * @param tree Tree
     * @param root 根
     */
    RootedTree(const Tree<Cost>& tree, int root) : Tree<Cost>(tree), root(root) {
        build();
    }
    /**
     * @brief Treeから根付き木を構築するコンストラクタ(ムーブ)
     *
     * @param tree Tree
     * @param root 根
     */
    RootedTree(Tree<Cost>&& tree, int root) : Tree<Cost>(std::move(tree)), root(root) {
        build();
    }
    
    /**
     * @brief ダブリングLCAが使える根付き木を得る
     * 
     * @return DoublingClimbTree<Cost> ダブリング済み根付き木のオブジェクト
     */
    [[nodiscard]]
    DoublingClimbTree<Cost> build_lca() const& {
        return DoublingClimbTree<Cost>(*this);
    }

    /**
     * @brief ダブリングLCAが使える根付き木を得る
     * 
     * @return DoublingClimbTree<Cost> ダブリング済み根付き木のオブジェクト
     */
    [[nodiscard]]
    DoublingClimbTree<Cost> build_lca() && {
        return DoublingClimbTree<Cost>(std::move(*this));
    }

private:
    void build() {
        int n = this->n;
        par.resize(n);
        par[root] = {root, -1, 0, -1};
        child.resize(n);
        depth.resize(n);
        unweighted_depth.resize(n);
        height.resize(n);
        unweighted_height.resize(n);
        sz.resize(n);
        std::vector<int> iter(n);
        std::stack<std::pair<int, int>> stk;
        stk.emplace(root, 0);
        while(!stk.empty()) {
            auto [u, cnt] = stk.top(); stk.pop();
            euler_tour.emplace_back(u, cnt);
            if(iter[u] == 0) preorder.push_back(u);
            while(iter[u] < (int)this->g[u].size() && this->g[u][iter[u]].dst == par[u]) iter[u]++;
            if(iter[u] == (int)this->g[u].size()) {
                postorder.push_back(u);
                sz[u] = 1;
                height[u] = 0;
                unweighted_height[u] = 0;
                for(auto& e : child[u]) {
                    sz[u] += sz[e.dst];
                    if(height[u] < height[e.dst] + e.cost) {
                        height[u] = height[e.dst] + e.cost;
                    }
                    if(unweighted_height[u] < unweighted_height[e.dst] + 1) {
                        unweighted_height[u] = unweighted_height[e.dst] + 1;
                    }
                }
            } else {
                auto& e = this->g[u][iter[u]++];
                par[e.dst] = {e.dst, u, e.cost, e.id};
                child[u].push_back(e);
                depth[e.dst] = depth[u] + e.cost;
                unweighted_depth[e.dst] = unweighted_depth[u] + 1;
                stk.emplace(u, cnt + 1);
                stk.emplace(e.dst, 0);
            }
        }
    }
};

/**
 * @brief 親頂点ダブリング済み根付き木
 * 
 * @tparam Cost = int 辺の重み(注: climbでは根の重みは考えない)
 */
template <typename Cost = int>
struct DoublingClimbTree : private RootedTree<Cost> {
    using Edge = typename RootedTree<Cost>::Edge;
    using RootedTree<Cost>::n;
    using RootedTree<Cost>::m;
    using RootedTree<Cost>::g;
    using RootedTree<Cost>::operator[];
    using RootedTree<Cost>::edges;
    using RootedTree<Cost>::shortest_path;
    using RootedTree<Cost>::root;
    using RootedTree<Cost>::par;
    using RootedTree<Cost>::child;
    using RootedTree<Cost>::depth;
    using RootedTree<Cost>::height;
    using RootedTree<Cost>::sz;
    using RootedTree<Cost>::preorder;
    using RootedTree<Cost>::postorder;
    using RootedTree<Cost>::euler_tour;

    int h; //!< ダブリングの回数
    std::vector<std::vector<int>> doubling_par; //!< jから(1<<i)頂点登った頂点(存在しない場合は-1)

    DoublingClimbTree() = delete;
    
    /**
     * @brief RootedTreeからダブリング済み根付き木を構築するコンストラクタ(コピー)
     *
     * @param tree RootedTree
     */
    DoublingClimbTree(const RootedTree<Cost>& tree) : RootedTree<Cost>(tree) {
        build();
    }

    /**
     * @brief RootedTreeからダブリング済み根付き木を構築するコンストラクタ(ムーブ)
     *
     * @param tree RootedTree
     */
    DoublingClimbTree(RootedTree<Cost>&& tree) : RootedTree<Cost>(std::move(tree)) {
        build();
    }
    
    /**
     * @brief 親の頂点のリストから根が0のダブリング済み根付き木を構築するコンストラクタ
     * 
     * @param par_ 頂点0以外の親の頂点のリスト
     * @param padding = -1 parの頂点番号をいくつずらすか
     */
    DoublingClimbTree(const std::vector<int>& par_, int padding = -1) : RootedTree<Cost>(par_, padding) {
        build();
    }
    /**
     * @brief Treeからダブリング済み根付き木を構築するコンストラクタ(コピー)
     *
     * @param tree Tree
     * @param root 根
     */
    DoublingClimbTree(const Tree<Cost>& tree, int root) : RootedTree<Cost>(tree, root) {
        build();
    }
    /**
     * @brief Treeからダブリング済み根付き木を構築するコンストラクタ(ムーブ)
     *
     * @param tree Tree
     * @param root 根
     */
    DoublingClimbTree(Tree<Cost>&& tree, int root) : RootedTree<Cost>(std::move(tree), root) {
        build();
    }
    
    /**
     * @brief 頂点uからk回を根の方向に遡った頂点
     * 
     * @param u 元の頂点
     * @param k 登る段数
     * @return int k回登った頂点
     */
    int climb(int u, int k) const {
        for(int i = 0; i < h; i++) {
            if(k >> i & 1) u = doubling_par[i][u];
            if(u == -1) return -1;
        }
        return u;
    }
    
    /**
     * @brief 最小共通祖先
     * 
     * @param u 頂点1
     * @param v 頂点2
     * @return int LCAの頂点番号
     */
    int lca(int u, int v) const {
        if(this->unweighted_depth[u] > this->unweighted_depth[v]) std::swap(u, v);
        v = climb(v, this->unweighted_depth[v] - this->unweighted_depth[u]);
        if(u == v) return u;
        for(int i = h - 1; i >= 0; i--) {
            int nu = doubling_par[i][u];
            int nv = doubling_par[i][v];
            if(nu == -1) continue;
            if(nu != nv) {
                u = nu;
                v = nv;
            }
        }
        return this->par[u];
    }
    
    /**
     * @brief 頂点間距離(重みなし)
     * 
     * @param u 頂点1
     * @param v 頂点2
     * @return int uとv間の最短経路の辺の本数
     */
    Cost dist(int u, int v) const {
        return this->depth[u] + this->depth[v] - this->depth[lca(u, v)] * 2;
    }
    
private:
    void build() {
        int n = this->n;
        h = 0;
        while((1 << h) < n) h++;
        doubling_par.assign(h, std::vector<int>(n, -1));
        for(int i = 0; i < n; i++) doubling_par[0][i] = this->par[i];
        for(int i = 0; i < h - 1; i++) {
            for(int j = 0; j < n; j++) {
                if(doubling_par[i][j] != -1) {
                    doubling_par[i+1][j] = doubling_par[i][doubling_par[i][j]];
                }
            }
        }
    }
};
#line 4 "test/yukicoder-898.test.cpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

int main() {
    int n; cin >> n;
    Tree<ll> tr(n); tr.read(0, true);
    DoublingClimbTree<ll> dct(tr, 0);
    int q; cin >> q;
    while(q--) {
        array<int, 3> xyz;
        for(int i = 0; i < 3; i++) cin >> xyz[i];
        int lca01 = dct.lca(xyz[0], xyz[1]);
        int lca02 = dct.lca(xyz[0], xyz[2]);
        int lca12 = dct.lca(xyz[1], xyz[2]);
        int center = lca01 ^ lca02 ^ lca12;
        cout << dct.dist(xyz[0], center) + dct.dist(xyz[1], center) + dct.dist(xyz[2], center) << endl;
    }
}