#995780 (C++23) No.901 K-ary εxtrεεmε

提出ソース
結果

問題	No.901 K-ary εxtrεεmε
コンテスト
ユーザー	cutmdo
提出日時	2024-07-18 22:44:34
言語	C++23 (gcc 13.3.0 + boost 1.87.0)
結果	AC
実行時間	376 ms / 3,000 ms
コード長	12,683 bytes
コンパイル時間	2,599 ms
コンパイル使用メモリ	171,120 KB
実行使用メモリ	54,676 KB
最終ジャッジ日時	2024-07-18 22:44:47
合計ジャッジ時間	12,539 ms
ジャッジサーバーID （参考情報）	judge5 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 1
other	AC * 29
権限があれば一括ダウンロードができます
ソースコード

raw source code
#define PROBLEM "https://yukicoder.me/problems/no/901"

#include <iostream>
#include <ranges>

#include <vector>
#include <set>
#include <unordered_map>
#include <queue>
#include <stack>

#include <vector>
#include <deque>

template<class Node = int, class Cost = long long>
class Graph {
    //using Node = int;
    //using Cost = long long;

    using Edge = std::pair<Node, Cost>;
    using Edges = std::vector<Edge>;

    const int m_n;
    std::vector<Edges> m_graph;

public:
    Graph(int n) :m_n(n), m_graph(n) {}

    auto addEdge(const Node& f, const Node& t, const Cost& c = 1) {
        m_graph[f].emplace_back(t, c);
    }
    auto addEdgeUndirected(const Node& f, const Node& t, const Cost& c = 1) {
        addEdge(f, t, c); addEdge(t, f, c);
    }
    auto getEdges(const Node& from)const {
        class EdgesRange {
            const typename Edges::const_iterator b, e;
        public:
            EdgesRange(const Edges& edges) :b(edges.begin()), e(edges.end()) {}
            auto begin()const { return b; }
            auto end()const { return e; }
        };
        return EdgesRange(m_graph[from]);
    }
    auto getEdges()const {
        std::deque<std::tuple<Node, Node, Cost>> edges;
        for(Node from = 0; from < m_n; ++from) for(const auto& [to, c] : getEdges(from)) {
            edges.emplace_back(from, to, c);
        }
        return edges;
    }
    auto getEdgesExcludeCost()const {
        std::deque<std::pair<Node, Node>> edges;
        for(Node from = 0; from < m_n; ++from) for(const auto& [to, _] : getEdges(from)) {
            edges.emplace_back(from, to);
        }
        return edges;
    }
    auto reverse()const {
        auto rev = Graph<Node, Cost>(m_n);
        for(const auto& [from, to, c] : getEdges()) {
            rev.addEdge(to, from, c);
        }
        return rev;
    }
    auto size()const { return m_n; };
};
template<class Node, class Cost>
class HeavyLightDecomposition {

    using GraphOrderd = std::unordered_map<Node, std::deque<Node>>;

    const Node m_n;
    const std::vector<Node> m_size;
    const GraphOrderd m_tree;
    const std::vector<Node> m_height;
    const std::vector<std::pair<Node, Node>> m_root_par;
    const std::vector<Node> m_ids;
    const std::vector<Node> m_order;
    const std::vector<Node> m_edge_ids;

    static auto constructGraph(const Graph<Node, Cost>& tree) {
        auto n = tree.size();
        std::deque<std::pair<Node, Node>> order;
        std::vector<Node> used(n);
        std::stack<std::pair<Node, Node>> stk;
        stk.emplace(0, -1); used[0] = true;
        while(!stk.empty()) {
            auto [f, p] = stk.top();
            order.emplace_front(f, p);
            stk.pop();
            for(const auto& [t, _] : tree.getEdges(f)) {
                if(used[t]) { continue;; }
                used[t] = true;
                stk.emplace(t, f);
            }
        }

        std::vector<Node> size(n, 1);
        GraphOrderd hld_tree;
        for(const auto& [f, p] : order) {
            Node size_sum = 1;
            Node size_max = 0;
            std::deque<Node> to_list;
            for(const auto& [t, _] : tree.getEdges(f)) {
                if(t == p) { continue; }
                if(size[t] > size_max) {
                    size_max = size[t];
                    to_list.emplace_back(t);
                } else {
                    to_list.emplace_front(t);
                }
                size_sum += size[t];
            }
            if(!to_list.empty()) {
                hld_tree.emplace(f, to_list);
            }
            size[f] = size_sum;
        }
        return hld_tree;
    }

    static auto constructSize(const Graph<Node, Cost>& tree) {
        auto n = tree.size();
        std::deque<std::pair<Node, Node>> order;
        std::vector<Node> used(n);
        std::stack<std::pair<Node, Node>> stk;
        stk.emplace(0, -1); used[0] = true;
        while(!stk.empty()) {
            auto [f, p] = stk.top();
            order.emplace_front(f, p);
            stk.pop();
            for(const auto& [t, _] : tree.getEdges(f)) {
                if(used[t]) { continue;; }
                used[t] = true;
                stk.emplace(t, f);
            }
        }

        std::vector<Node> size(n, 1);
        for(const auto& [f, p] : order) {
            Node size_sum = 1;
            for(const auto& [t, _] : tree.getEdges(f)) {
                if(t == p) { continue; }
                size_sum += size[t];
            }
            size[f] = size_sum;
        }
        return size;

    }

    static auto constructRootPar(Node n, const GraphOrderd& tree) {
        std::vector<std::pair<Node, Node>> root_par(n);
        std::stack<std::tuple<Node, Node, Node>> stk;
        stk.emplace(0, 0, -1);
        while(!stk.empty()) {
            auto [f, root, par] = stk.top();
            stk.pop();

            if(tree.find(f) == tree.end()) { root_par[f] = {root,par}; continue; }
            auto itr = tree.at(f).rbegin();
            stk.emplace(*itr, root, par);
            root_par[f] = {root,par};
            for(++itr; itr != tree.at(f).rend(); ++itr) {
                stk.emplace(*itr, *itr, f);
            }
        }
        return root_par;
    }
    static auto constructHeight(Node n, const GraphOrderd& tree) {
        std::vector<Node> height(n);
        std::queue<Node> q;
        q.emplace(0);
        while(!q.empty()) {
            auto f = q.front();
            q.pop();
            if(tree.find(f) == tree.end()) { continue; }
            for(const auto& t : tree.at(f)) {
                height[t] = height[f] + 1;
                q.emplace(t);
            }
        }
        return height;
    }

    auto constructIds() const {
        std::vector<Node> ids(m_n);
        Node val = 0;
        std::stack<Node> stk;
        stk.emplace(0);
        while(!stk.empty()) {
            auto f = stk.top();
            stk.pop();
            ids[f] = val; ++val;
            if(m_tree.find(f) == m_tree.end()) { continue; }
            for(const auto& t : m_tree.at(f)) { stk.emplace(t); }
        }
        return ids;
    }

    auto constructOrder()const {
        std::vector<Node> order(m_n);
        for(int i = 0; i < m_n; ++i) { order[m_ids[i]] = i; }
        return order;
    }
    /*
     * 辺をnodeとして拡張した場合の辺nodeだけIDを振る
     * (1) - (2)
     * (1) - (e) - (2)
     * [-1, -1, 0]
     */
    auto constructEdgeIds() const {
        Node edge_size = (m_n >> 1);
        std::vector<Node> edge_ids(m_n, -1);
        Node val = 0;
        std::stack<Node> stk;
        stk.emplace(0);
        while(!stk.empty()) {
            auto f = stk.top();
            stk.pop();
            if(f > edge_size) { edge_ids[f] = val; ++val; }
            if(m_tree.find(f) == m_tree.end()) { continue; }
            for(const auto& t : m_tree.at(f)) { stk.emplace(t); }
        }
        return edge_ids;
    }

public:

    HeavyLightDecomposition(const Graph<Node, Cost>& tree) :
        m_n(tree.size()),
        m_size(constructSize(tree)),
        m_tree(constructGraph(tree)),
        m_root_par(constructRootPar(m_n, m_tree)),
        m_height(constructHeight(m_n, m_tree)),
        m_ids(constructIds()),
        m_order(constructOrder()),
        m_edge_ids(constructEdgeIds()) {
    }

    auto getId(Node i)const { return m_ids[i]; }
    auto getEdgeId(Node i)const { return m_edge_ids[i]; }
    auto getOrder(Node i)const { return m_order[i]; }

    auto lca(Node f, Node t)const {
        do {
            auto [fr, fp] = m_root_par[f];
            auto [tr, tp] = m_root_par[t];
            if(fr == tr) { break; }
            auto fph = (fp > -1) ? m_height[fp] : -1;
            auto tph = (tp > -1) ? m_height[tp] : -1;
            if(fph < tph) { t = tp; } else { f = fp; }
        } while(true);
        return (m_height[f] < m_height[t]) ? f : t;
    }

    auto range(Node f, Node t)const {
        std::deque<std::pair<Node, Node>> ret;
        auto add = [&](Node f, Node t) {
            auto l = std::min(m_ids[f], m_ids[t]);
            auto r = std::max(m_ids[f], m_ids[t]);
            ret.emplace_back(l, r);
        };
        do {
            auto [fr, fp] = m_root_par[f];
            auto [tr, tp] = m_root_par[t];
            if(fr == tr) { add(f, t); break; }
            auto fph = (fp > -1) ? m_height[fp] : -1;
            auto tph = (tp > -1) ? m_height[tp] : -1;
            if(fph < tph) { add(t, tr); t = tp; } else { add(f, fr); f = fp; }
        } while(true);
        return ret;
    }

    auto rangeEdge(Node f, Node t)const {
        Node edge_size = (m_n >> 1);
        std::deque<std::pair<Node, Node>> ret;
        auto add = [&](Node f, Node t) {
            auto l = std::min(m_ids[f], m_ids[t]);
            auto r = std::max(m_ids[f], m_ids[t]);
            if(m_order[l] <= edge_size) { ++l; }
            if(m_order[r] <= edge_size) { --r; }
            if(l > r) { return; }
            auto edge_l = m_edge_ids[m_order[l]];
            auto edge_r = m_edge_ids[m_order[r]];
            ret.emplace_back(edge_l, edge_r);
        };
        do {
            auto [fr, fp] = m_root_par[f];
            auto [tr, tp] = m_root_par[t];
            if(fr == tr) { add(f, t); break; }
            auto fph = (fp > -1) ? m_height[fp] : -1;
            auto tph = (tp > -1) ? m_height[tp] : -1;
            if(fph < tph) { add(t, tr); t = tp; } else { add(f, fr); f = fp; }
        } while(true);
        return ret;
    }

    auto rangeSubTree(Node f)const {
        return std::pair<Node, Node>{
            m_ids[f], m_ids[f] + m_size[f] - 1
        };
    }
};

template<class Node, class Cost>
class AuxiliaryTree {

    // 定数倍高速化のため破壊的
    std::vector<int> compres_map;

    const std::vector<Cost> depth_cost;
    const HeavyLightDecomposition<Node, Cost> hld;

    auto construct_depth(const Graph<Node, Cost>& tree)const {
        std::vector<Cost> depth_cost(tree.size());
        std::vector<int> used(tree.size());
        auto dfs = [&](auto&& self, Node from) -> void {
            used[from] = true;
            for(const auto& [to, c] : tree.getEdges(from)) if(!used[to]) {
                depth_cost[to] = depth_cost[from] + c;
                self(self, to);
            }
        };
        dfs(dfs, 0);
        return depth_cost;
    }

public:
    AuxiliaryTree(const Graph<Node, Cost>& tree) :
        compres_map(tree.size()),
        depth_cost(construct_depth(tree)),
        hld(tree) {
    }

    auto compression(const std::vector<int>& nodes) {
        auto compare = [&](int a, int b) {return hld.getId(a) < hld.getId(b); };

        // 元の頂点集合
        auto nodes_set = std::set<int, decltype(compare)>(nodes.begin(), nodes.end(), compare);
        auto nodes_set_with_lca = nodes_set;

        // pre orderでの全ての隣接nodeのLCAを求める
        for(auto itr = nodes_set_with_lca.begin(); std::next(itr) != nodes_set_with_lca.end(); ++itr) {
            nodes_set_with_lca.emplace(hld.lca(*itr, *std::next(itr)));
        }

        // 座標圧縮
        int at_size = nodes_set_with_lca.size();
        for(int i = 0; auto x : nodes_set_with_lca) { compres_map[x] = i; ++i; }

        // LCAを含めた全てのnodeで子孫関係を保って辺を張る
        std::stack<int> stk;
        Graph<Node, Cost> auxiliary_tree(at_size);
        for(auto nd : nodes_set_with_lca) {
            while(!stk.empty() && hld.lca(stk.top(), nd) != stk.top()) {
                stk.pop();
            }
            if(!stk.empty()) {
                auto f = compres_map[stk.top()];
                auto t = compres_map[nd];
                auto c = depth_cost[stk.top()] + depth_cost[nd] - depth_cost[hld.lca(stk.top(), nd)] * 2;
                auxiliary_tree.addEdgeUndirected(f, t, c);
            }
            stk.emplace(nd);
        }
        return auxiliary_tree;
    }
};

int main() {
    std::cin.tie(0);
    std::ios::sync_with_stdio(0);

    int n;
    std::cin >> n;
    Graph tree(n);
    for(auto _ : std::views::iota(0, n - 1)) {
        int f, t, c;
        std::cin >> f >> t >> c;
        tree.addEdgeUndirected(f, t, c);
    }

    auto at = AuxiliaryTree(tree);

    int q;
    std::cin >> q;
    for([[maybe_unused]] auto _ : std::views::iota(0, q)) {
        int k;
        std::cin >> k;
        std::vector<int> v(k);
        for(auto i : std::views::iota(0, k)) { std::cin >> v[i]; }

        auto comp_tree = at.compression(v);

        long long ans = 0;
        for(const auto& [f, t, c] : comp_tree.getEdges()) if(f < t) {
            ans += c;
        }

        std::cout << ans << "\n";
    }
}