cp-algo/graph/mst.hpp

• View this file on GitHub
• Last update: 2025-12-12 22:29:00+01:00
• Include: #include "cp-algo/graph/mst.hpp"

Depends on

• cp-algo/graph/base.hpp
• cp-algo/graph/concepts.hpp
• cp-algo/graph/edge_types.hpp
• cp-algo/structures/dsu.hpp
• cp-algo/structures/stack_union.hpp
• cp-algo/util/big_alloc.hpp
• cp-algo/util/bit.hpp
• cp-algo/util/simd.hpp
• cp-algo/util/sort.hpp

Verified with

• Minimum Spanning Tree (verify/graph/mst.test.cpp)

Code

#ifndef CP_ALGO_GRAPH_MST_HPP
#define CP_ALGO_GRAPH_MST_HPP
#include "base.hpp"
#include "../structures/dsu.hpp"
#include "../util/sort.hpp"
#include <algorithm>
namespace cp_algo::graph {
    template<weighted_undirected_graph_type graph>
    std::pair<int64_t, big_vector<edge_index>> mst(graph const& g) {
        struct edge {
            edge_index idx;
            node_index v;
        };
        big_vector<edge> edges;
        for(auto v: g.nodes()) {
            for(auto e: g.outgoing(v)) {
                if (v < g.edge(e).traverse(v)) {
                    edges.push_back({e, v});
                }
            }
        }
        radix_sort(edges, [&](auto e) {
            return g.edge(e.idx).w;
        });
        structures::dsu me(g.n());
        int64_t total = 0;
        big_vector<edge_index> mst;
        for(auto [idx, v]: edges) {
            if(me.uni(v, g.edge(idx).traverse(v))) {
                total += g.edge(idx).w;
                mst.push_back(idx);
            }
        }
        return {total, mst};
    }
}
#endif // CP_ALGO_GRAPH_MST_HPP

#line 1 "cp-algo/graph/mst.hpp"


#line 1 "cp-algo/graph/base.hpp"


#line 1 "cp-algo/graph/edge_types.hpp"


#include <iostream>
#include <cstdint>
namespace cp_algo::graph {
    using node_index = int;
    struct edge_base {
        int xor_nodes;

        edge_base() {}
        edge_base(node_index from, node_index to): xor_nodes(from ^ to) {}

        // Given one endpoint, return the other
        node_index traverse(node_index from) const {
            return xor_nodes ^ from;
        }

        static auto read(node_index v0 = 0) {
            node_index u, v;
            std::cin >> u >> v;
            u -= v0;
            v -= v0;
            return std::pair{u, edge_base(u, v)};
        }
    };

    struct weighted_edge: edge_base {
        int64_t w;

        weighted_edge() {}
        weighted_edge(node_index from, node_index to, int64_t w): edge_base(from, to), w(w) {}

        static auto read(node_index v0 = 0) {
            auto [u, e] = edge_base::read(v0);
            int64_t w;
            std::cin >> w;
            return std::pair{u, weighted_edge(u, e.traverse(u), w)};
        }
    };

    template<typename edge>
    concept edge_type = std::is_base_of_v<edge_base, edge>;
    template<typename edge>
    concept weighted_edge_type = std::is_base_of_v<weighted_edge, edge>;
}

#line 1 "cp-algo/graph/concepts.hpp"


#line 4 "cp-algo/graph/concepts.hpp"
#include <type_traits>

namespace cp_algo::graph {
    // Shared graph mode enum for all graph headers
    enum graph_mode { directed, undirected };
    // Traits: true for types that expose `edge_t` and static `mode`
    template<typename T, typename = void>
    struct graph_traits : std::false_type {};

    template<typename T>
    struct graph_traits<T, std::void_t<typename T::edge_t, decltype(T::mode)>> : std::true_type {
        using edge_t = typename T::edge_t;
        static constexpr auto mode = T::mode;
        static constexpr bool is_directed = mode == directed;
        static constexpr bool is_undirected = mode == undirected;
        static constexpr bool is_weighted = weighted_edge_type<edge_t>;
    };

    // Concepts
    template<typename G>
    concept graph_type = graph_traits<G>::value;

    template<typename G>
    concept digraph_type = graph_type<G> && graph_traits<G>::is_directed;

    template<typename G>
    concept undirected_graph_type = graph_type<G> && graph_traits<G>::is_undirected;

    template<typename G>
    concept weighted_graph_type = graph_type<G> && graph_traits<G>::is_weighted;

    template<typename G>
    concept weighted_digraph_type = digraph_type<G> && graph_traits<G>::is_weighted;

    template<typename G>
    concept weighted_undirected_graph_type = undirected_graph_type<G> && graph_traits<G>::is_weighted;
}

#line 1 "cp-algo/structures/stack_union.hpp"


#line 1 "cp-algo/util/big_alloc.hpp"



#include <set>
#include <map>
#include <deque>
#include <stack>
#include <queue>
#include <vector>
#include <string>
#include <cstddef>
#line 13 "cp-algo/util/big_alloc.hpp"
#include <generator>
#include <forward_list>

// Single macro to detect POSIX platforms (Linux, Unix, macOS)
#if defined(__linux__) || defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#  define CP_ALGO_USE_MMAP 1
#  include <sys/mman.h>
#else
#  define CP_ALGO_USE_MMAP 0
#endif

namespace cp_algo {
    template <typename T, size_t Align = 32>
    class big_alloc {
        static_assert( Align >= alignof(void*), "Align must be at least pointer-size");
        static_assert(std::popcount(Align) == 1, "Align must be a power of two");
    public:
        using value_type = T;
        template <class U> struct rebind { using other = big_alloc<U, Align>; };
        constexpr bool operator==(const big_alloc&) const = default;
        constexpr bool operator!=(const big_alloc&) const = default;

        big_alloc() noexcept = default;
        template <typename U, std::size_t A>
        big_alloc(const big_alloc<U, A>&) noexcept {}

        [[nodiscard]] T* allocate(std::size_t n) {
            std::size_t padded = round_up(n * sizeof(T));
            std::size_t align = std::max<std::size_t>(alignof(T),  Align);
#if CP_ALGO_USE_MMAP
            if (padded >= MEGABYTE) {
                void* raw = mmap(nullptr, padded,
                                PROT_READ | PROT_WRITE,
                                MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
                madvise(raw, padded, MADV_HUGEPAGE);
                madvise(raw, padded, MADV_POPULATE_WRITE);
                return static_cast<T*>(raw);
            }
#endif
            return static_cast<T*>(::operator new(padded, std::align_val_t(align)));
        }

        void deallocate(T* p, std::size_t n) noexcept {
            if (!p) return;
            std::size_t padded = round_up(n * sizeof(T));
            std::size_t align  = std::max<std::size_t>(alignof(T),  Align);
    #if CP_ALGO_USE_MMAP
            if (padded >= MEGABYTE) { munmap(p, padded); return; }
    #endif
            ::operator delete(p, padded, std::align_val_t(align));
        }

    private:
        static constexpr std::size_t MEGABYTE = 1 << 20;
        static constexpr std::size_t round_up(std::size_t x) noexcept {
            return (x + Align - 1) / Align * Align;
        }
    };

    template<typename T> using big_vector = std::vector<T, big_alloc<T>>;
    template<typename T> using big_basic_string = std::basic_string<T, std::char_traits<T>, big_alloc<T>>;
    template<typename T> using big_deque = std::deque<T, big_alloc<T>>;
    template<typename T> using big_stack = std::stack<T, big_deque<T>>;
    template<typename T> using big_queue = std::queue<T, big_deque<T>>;
    template<typename T> using big_priority_queue = std::priority_queue<T, big_vector<T>>;
    template<typename T> using big_forward_list = std::forward_list<T, big_alloc<T>>;
    using big_string = big_basic_string<char>;

    template<typename Key, typename Value, typename Compare = std::less<Key>>
    using big_map = std::map<Key, Value, Compare, big_alloc<std::pair<const Key, Value>>>;
    template<typename T, typename Compare = std::less<T>>
    using big_multiset = std::multiset<T, Compare, big_alloc<T>>;
    template<typename T, typename Compare = std::less<T>>
    using big_set = std::set<T, Compare, big_alloc<T>>;
    template<typename Ref, typename V = void>

    using big_generator = std::generator<Ref, V, big_alloc<std::byte>>;
}

// Deduction guide to make elements_of with big_generator default to big_alloc
namespace std::ranges {
    template<typename Ref, typename V>
    elements_of(cp_algo::big_generator<Ref, V>&&) -> elements_of<cp_algo::big_generator<Ref, V>&&, cp_algo::big_alloc<std::byte>>;
}


#line 5 "cp-algo/structures/stack_union.hpp"
#include <iterator>
#include <ranges>
namespace cp_algo::structures {
    template<class datatype>
    struct stack_union {
        stack_union(int n = 0): head(n), next(1), data(1) {}

        void push(int v, datatype const& vdata) {
            next.push_back(head[v]);
            head[v] = (int)std::size(next) - 1;
            data.push_back(vdata);
        }
        template<typename... Args>
        void emplace(int v, Args&&... vdata) {
            next.push_back(head[v]);
            head[v] = (int)std::size(next) - 1;
            data.emplace_back(std::forward<Args>(vdata)...);
        }

        void reserve(int m) {
            data.reserve(m);
            next.reserve(m);
        }

        size_t size() const {return std::size(head);}
        size_t nodes() const {return std::size(data);}

        template<typename Su>
        struct _iterator {
            using value_type = std::conditional_t<std::is_const_v<Su>, const datatype, datatype>;
            using difference_type = std::ptrdiff_t;

            Su* su = nullptr;
            int sv = 0;

            value_type& operator*() const { return su->data[sv]; }
            _iterator& operator++() { 
                sv = su->next[sv];
                return *this; 
            }
            _iterator operator++(int) { auto tmp = *this; ++*this; return tmp; }
            friend bool operator==(_iterator const& it, std::default_sentinel_t) { 
                return it.sv == 0;
            }
        };

        using iterator = _iterator<stack_union<datatype>>;
        using const_iterator = _iterator<const stack_union<datatype>>;

        auto operator[](this auto&& self, int v) {
            using Iter = _iterator<std::remove_reference_t<decltype(self)>>;
            return std::ranges::subrange(Iter{&self, self.head[v]}, std::default_sentinel);
        }

        big_vector<int> head, next;
        big_vector<datatype> data;
    };
}

#line 7 "cp-algo/graph/base.hpp"
namespace cp_algo::graph {
    using edge_index = int;
    template<edge_type _edge_t = edge_base, graph_mode _mode = undirected>
    struct graph {
        using edge_t = _edge_t;
        static constexpr auto mode = _mode;
        using incidence_list = structures::stack_union<edge_index>;
        graph(int n, int v0 = 0): v0(v0), adj(n) {}

        graph transpose() const {
            static_assert(mode == directed, "transpose is only defined for directed graphs");
            graph<edge_t, mode> gt(n(), v0);
            for(auto v: nodes()) {
                for(auto e: outgoing(v)) {
                    gt.add_edge(edge(e).traverse(v), edge(e));
                }
            }
            return gt;
        }
        edge_index add_edge(node_index u, edge_t e) {
            edge_index idx = (edge_index)size(E);
            E.push_back(e);
            adj.push(u, idx);
            if constexpr (mode == undirected) {
                adj.push(e.traverse(u), idx);
            }
            return idx;
        }
        edge_index add_edge(node_index u, auto... Args) {
            return add_edge(u, edge_t(u, Args...));
        }
        void read_edges(node_index m) {
            adj.reserve(mode == undirected ? 2 * m : m);
            for(edge_index i = 0; i < m; i++) {
                auto [u, e] = edge_t::read(v0);
                add_edge(u, e);
            }
        }
        auto outgoing(node_index v) const {return adj[v];}
        auto edges() const {return E | std::views::all;}
        auto nodes() const {return std::views::iota(node_index(0), n());}
        auto edge_indices() const {return std::views::iota(edge_index(0), m());}
        auto&& incidence_lists(this auto&& self) {return self.adj;}
        auto&& edge(this auto&& self, edge_index e) {return self.E[e];}
        node_index n() const {return (node_index)incidence_lists().size();}
        edge_index m() const {return (edge_index)edges().size();}
    private:
        node_index v0;
        big_vector<edge_t> E;
        incidence_list adj;
    };
    // aliases for most standard cases
    template<edge_type edge_t = edge_base>
    using digraph = graph<edge_t, directed>;
    template<weighted_edge_type edge_t = weighted_edge, graph_mode mode = undirected>
    using weighted_graph = graph<edge_t, mode>;
    template<weighted_edge_type edge_t = weighted_edge>
    using weighted_digraph = digraph<edge_t>;
}

#line 1 "cp-algo/structures/dsu.hpp"


#line 4 "cp-algo/structures/dsu.hpp"
#include <numeric>
#line 6 "cp-algo/structures/dsu.hpp"
namespace cp_algo::structures {
    struct disjoint_set_union {
        disjoint_set_union (int n): par(n) {
            std::iota(begin(par), end(par), 0);
        }
        int get(int v) {
            return v == par[v] ? v : par[v] = get(par[v]);
        }
        bool uni(int a, int b) {
            a = get(a);
            b = get(b);
            par[a] = b;
            return a != b;
        }
    private:
        big_vector<int> par;
    };
    using dsu = disjoint_set_union;
}

#line 1 "cp-algo/util/sort.hpp"


#line 1 "cp-algo/util/bit.hpp"


#line 1 "cp-algo/util/simd.hpp"


#include <experimental/simd>
#line 6 "cp-algo/util/simd.hpp"
#include <memory>

#if defined(__x86_64__) && !defined(CP_ALGO_DISABLE_AVX2)
#define CP_ALGO_SIMD_AVX2_TARGET _Pragma("GCC target(\"avx2\")")
#else
#define CP_ALGO_SIMD_AVX2_TARGET
#endif

#define CP_ALGO_SIMD_PRAGMA_PUSH \
    _Pragma("GCC push_options") \
    CP_ALGO_SIMD_AVX2_TARGET

CP_ALGO_SIMD_PRAGMA_PUSH
namespace cp_algo {
    template<typename T, size_t len>
    using simd [[gnu::vector_size(len * sizeof(T))]] = T;
    using i64x4 = simd<int64_t, 4>;
    using u64x4 = simd<uint64_t, 4>;
    using u32x8 = simd<uint32_t, 8>;
    using i32x4 = simd<int32_t, 4>;
    using u32x4 = simd<uint32_t, 4>;
    using i16x4 = simd<int16_t, 4>;
    using u8x32 = simd<uint8_t, 32>;
    using dx4 = simd<double, 4>;

    dx4 abs(dx4 a) {
        return dx4{
            std::abs(a[0]),
            std::abs(a[1]),
            std::abs(a[2]),
            std::abs(a[3])
        };
    }

    // https://stackoverflow.com/a/77376595
    // works for ints in (-2^51, 2^51)
    static constexpr dx4 magic = dx4() + (3ULL << 51);
    inline i64x4 lround(dx4 x) {
        return i64x4(x + magic) - i64x4(magic);
    }
    inline dx4 to_double(i64x4 x) {
        return dx4(x + i64x4(magic)) - magic;
    }

    inline dx4 round(dx4 a) {
        return dx4{
            std::nearbyint(a[0]),
            std::nearbyint(a[1]),
            std::nearbyint(a[2]),
            std::nearbyint(a[3])
        };
    }

    inline u64x4 low32(u64x4 x) {
        return x & uint32_t(-1);
    }
    inline auto swap_bytes(auto x) {
        return decltype(x)(__builtin_shufflevector(u32x8(x), u32x8(x), 1, 0, 3, 2, 5, 4, 7, 6));
    }
    inline u64x4 montgomery_reduce(u64x4 x, uint32_t mod, uint32_t imod) {
#ifdef __AVX2__
        auto x_ninv = u64x4(_mm256_mul_epu32(__m256i(x), __m256i() + imod));
        x += u64x4(_mm256_mul_epu32(__m256i(x_ninv), __m256i() + mod));
#else
        auto x_ninv = u64x4(u32x8(low32(x)) * imod);
        x += x_ninv * uint64_t(mod);
#endif
        return swap_bytes(x);
    }

    inline u64x4 montgomery_mul(u64x4 x, u64x4 y, uint32_t mod, uint32_t imod) {
#ifdef __AVX2__
        return montgomery_reduce(u64x4(_mm256_mul_epu32(__m256i(x), __m256i(y))), mod, imod);
#else
        return montgomery_reduce(x * y, mod, imod);
#endif
    }
    inline u32x8 montgomery_mul(u32x8 x, u32x8 y, uint32_t mod, uint32_t imod) {
        return u32x8(montgomery_mul(u64x4(x), u64x4(y), mod, imod)) |
               u32x8(swap_bytes(montgomery_mul(u64x4(swap_bytes(x)), u64x4(swap_bytes(y)), mod, imod)));
    }
    inline dx4 rotate_right(dx4 x) {
        static constexpr u64x4 shuffler = {3, 0, 1, 2};
        return __builtin_shuffle(x, shuffler);
    }

    template<std::size_t Align = 32>
    inline bool is_aligned(const auto* p) noexcept {
        return (reinterpret_cast<std::uintptr_t>(p) % Align) == 0;
    }

    template<class Target>
    inline Target& vector_cast(auto &&p) {
        return *reinterpret_cast<Target*>(std::assume_aligned<alignof(Target)>(&p));
    }
}
#pragma GCC pop_options

#line 6 "cp-algo/util/bit.hpp"
#include <array>
#include <bit>

#if defined(__x86_64__) && !defined(CP_ALGO_DISABLE_AVX2)
#define CP_ALGO_BIT_OPS_TARGET _Pragma("GCC target(\"avx2,bmi,bmi2,lzcnt,popcnt\")")
#else
#define CP_ALGO_BIT_OPS_TARGET _Pragma("GCC target(\"bmi,bmi2,lzcnt,popcnt\")")
#endif

#define CP_ALGO_BIT_PRAGMA_PUSH \
    _Pragma("GCC push_options") \
    CP_ALGO_BIT_OPS_TARGET

CP_ALGO_BIT_PRAGMA_PUSH
namespace cp_algo {
    template<typename Uint>
    constexpr size_t bit_width = sizeof(Uint) * 8;

    // n < 64
    uint64_t mask(size_t n) {
        return (1ULL << n) - 1;
    }
    size_t order_of_bit(auto x, size_t k) {
        return k ? std::popcount(x << (bit_width<decltype(x)> - k)) : 0;
    }
    inline size_t kth_set_bit(uint64_t x, size_t k) {
        return std::countr_zero(_pdep_u64(1ULL << k, x));
    }
    template<int fl = 0>
    void with_bit_floor(size_t n, auto &&callback) {
        if constexpr (fl >= 63) {
            return;
        } else if (n >> (fl + 1)) {
            with_bit_floor<fl + 1>(n, callback);
        } else {
            callback.template operator()<1ULL << fl>();
        }
    }
    void with_bit_ceil(size_t n, auto &&callback) {
        with_bit_floor(n, [&]<size_t N>() {
            if(N == n) {
                callback.template operator()<N>();
            } else {
                callback.template operator()<N << 1>();
            }
        });
    }

    inline uint32_t read_bits(char const* p) {
        return _mm256_movemask_epi8(__m256i(vector_cast<u8x32 const>(p[0]) + (127 - '0')));
    }
    inline uint64_t read_bits64(char const* p) {
        return read_bits(p) | (uint64_t(read_bits(p + 32)) << 32);
    }

    inline void write_bits(char *p, uint32_t bits) {
        static constexpr u8x32 shuffler = {
            0, 0, 0, 0, 0, 0, 0, 0,
            1, 1, 1, 1, 1, 1, 1, 1,
            2, 2, 2, 2, 2, 2, 2, 2,
            3, 3, 3, 3, 3, 3, 3, 3
        };
        auto shuffled = u8x32(_mm256_shuffle_epi8(__m256i() + bits, __m256i(shuffler)));
        static constexpr u8x32 mask = {
            1, 2, 4, 8, 16, 32, 64, 128,
            1, 2, 4, 8, 16, 32, 64, 128,
            1, 2, 4, 8, 16, 32, 64, 128,
            1, 2, 4, 8, 16, 32, 64, 128
        };
        for(int z = 0; z < 32; z++) {
            p[z] = shuffled[z] & mask[z] ? '1' : '0';
        }
    }
    inline void write_bits64(char *p, uint64_t bits) {
        write_bits(p, uint32_t(bits));
        write_bits(p + 32, uint32_t(bits >> 32));
    }
}
#pragma GCC pop_options

#line 4 "cp-algo/util/sort.hpp"
#include <algorithm>
#line 8 "cp-algo/util/sort.hpp"
namespace cp_algo {
    template<size_t maxc>
    void count_sort(auto &a, auto &&proj) {
        std::array<int, maxc> cnt = {};
        for(auto &x: a) {
            cnt[proj(x)]++;
        }
        std::partial_sum(begin(cnt), end(cnt), begin(cnt));
        auto res = a;
        for(auto const& it: a | std::views::reverse) {
            res[--cnt[proj(it)]] = it;
        }
        a = std::move(res);
    }
    template<size_t maxc>
    void count_sort(auto &a) {
        count_sort<maxc>(a, std::identity{});
    }

    void radix_sort(auto &a, auto &&proj) {
        if(empty(a)) {
            return;
        }
        auto [mn, mx] = std::ranges::minmax(a, {}, proj);
        with_bit_floor<1>(size(a), [&]<size_t floor>() {
            constexpr int base = std::min<size_t>(floor, 1 << 16);
            for(int64_t i = 1; i <= std::invoke(proj, mx) - std::invoke(proj, mn); i *= base) {
                count_sort<base>(a, [&](auto const& x) {
                    return (std::invoke(proj, x) - std::invoke(proj, mn)) / i % base;
                });
            }
        });
    }
    void radix_sort(auto &a) {
        radix_sort(a, std::identity{});
    }
}

#line 7 "cp-algo/graph/mst.hpp"
namespace cp_algo::graph {
    template<weighted_undirected_graph_type graph>
    std::pair<int64_t, big_vector<edge_index>> mst(graph const& g) {
        struct edge {
            edge_index idx;
            node_index v;
        };
        big_vector<edge> edges;
        for(auto v: g.nodes()) {
            for(auto e: g.outgoing(v)) {
                if (v < g.edge(e).traverse(v)) {
                    edges.push_back({e, v});
                }
            }
        }
        radix_sort(edges, [&](auto e) {
            return g.edge(e.idx).w;
        });
        structures::dsu me(g.n());
        int64_t total = 0;
        big_vector<edge_index> mst;
        for(auto [idx, v]: edges) {
            if(me.uni(v, g.edge(idx).traverse(v))) {
                total += g.edge(idx).w;
                mst.push_back(idx);
            }
        }
        return {total, mst};
    }
}

#ifndef CP_ALGO_GRAPH_MST_HPP
#define CP_ALGO_GRAPH_MST_HPP
#include "base.hpp"
#include "../structures/dsu.hpp"
#include "../util/sort.hpp"
#include <algorithm>
namespace cp_algo::graph{template<weighted_undirected_graph_type graph>std::pair<int64_t,big_vector<edge_index>>mst(graph const&g){struct edge{edge_index idx;node_index v;};big_vector<edge>edges;for(auto v:g.nodes()){for(auto e:g.outgoing(v)){if(v<g.edge(e).traverse(v)){edges.push_back({e,v});}}}radix_sort(edges,[&](auto e){return g.edge(e.idx).w;});structures::dsu me(g.n());int64_t total=0;big_vector<edge_index>mst;for(auto[idx,v]:edges){if(me.uni(v,g.edge(idx).traverse(v))){total+=g.edge(idx).w;mst.push_back(idx);}}return{total,mst};}}
#endif

#line 1 "cp-algo/graph/mst.hpp"
#line 1 "cp-algo/graph/base.hpp"
#line 1 "cp-algo/graph/edge_types.hpp"
#include <iostream>
#include <cstdint>
namespace cp_algo::graph{using node_index=int;struct edge_base{int xor_nodes;edge_base(){}edge_base(node_index from,node_index to):xor_nodes(from^to){}node_index traverse(node_index from)const{return xor_nodes^from;}static auto read(node_index v0=0){node_index u,v;std::cin>>u>>v;u-=v0;v-=v0;return std::pair{u,edge_base(u,v)};}};struct weighted_edge:edge_base{int64_t w;weighted_edge(){}weighted_edge(node_index from,node_index to,int64_t w):edge_base(from,to),w(w){}static auto read(node_index v0=0){auto[u,e]=edge_base::read(v0);int64_t w;std::cin>>w;return std::pair{u,weighted_edge(u,e.traverse(u),w)};}};template<typename edge>concept edge_type=std::is_base_of_v<edge_base,edge>;template<typename edge>concept weighted_edge_type=std::is_base_of_v<weighted_edge,edge>;}
#line 1 "cp-algo/graph/concepts.hpp"
#line 4 "cp-algo/graph/concepts.hpp"
#include <type_traits>
namespace cp_algo::graph{enum graph_mode{directed,undirected};template<typename T,typename=void>struct graph_traits:std::false_type{};template<typename T>struct graph_traits<T,std::void_t<typename T::edge_t,decltype(T::mode)>>:std::true_type{using edge_t=typename T::edge_t;static constexpr auto mode=T::mode;static constexpr bool is_directed=mode==directed;static constexpr bool is_undirected=mode==undirected;static constexpr bool is_weighted=weighted_edge_type<edge_t>;};template<typename G>concept graph_type=graph_traits<G>::value;template<typename G>concept digraph_type=graph_type<G>&&graph_traits<G>::is_directed;template<typename G>concept undirected_graph_type=graph_type<G>&&graph_traits<G>::is_undirected;template<typename G>concept weighted_graph_type=graph_type<G>&&graph_traits<G>::is_weighted;template<typename G>concept weighted_digraph_type=digraph_type<G>&&graph_traits<G>::is_weighted;template<typename G>concept weighted_undirected_graph_type=undirected_graph_type<G>&&graph_traits<G>::is_weighted;}
#line 1 "cp-algo/structures/stack_union.hpp"
#line 1 "cp-algo/util/big_alloc.hpp"
#include <set>
#include <map>
#include <deque>
#include <stack>
#include <queue>
#include <vector>
#include <string>
#include <cstddef>
#line 13 "cp-algo/util/big_alloc.hpp"
#include <generator>
#include <forward_list>
#if defined(__linux__) || defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#  define CP_ALGO_USE_MMAP 1
#  include <sys/mman.h>
#else
#  define CP_ALGO_USE_MMAP 0
#endif
namespace cp_algo{template<typename T,size_t Align=32>class big_alloc{static_assert(Align>=alignof(void*),"Align must be at least pointer-size");static_assert(std::popcount(Align)==1,"Align must be a power of two");public:using value_type=T;template<class U>struct rebind{using other=big_alloc<U,Align>;};constexpr bool operator==(const big_alloc&)const=default;constexpr bool operator!=(const big_alloc&)const=default;big_alloc()noexcept=default;template<typename U,std::size_t A>big_alloc(const big_alloc<U,A>&)noexcept{}[[nodiscard]]T*allocate(std::size_t n){std::size_t padded=round_up(n*sizeof(T));std::size_t align=std::max<std::size_t>(alignof(T),Align);
#if CP_ALGO_USE_MMAP
if(padded>=MEGABYTE){void*raw=mmap(nullptr,padded,PROT_READ|PROT_WRITE,MAP_PRIVATE|MAP_ANONYMOUS,-1,0);madvise(raw,padded,MADV_HUGEPAGE);madvise(raw,padded,MADV_POPULATE_WRITE);return static_cast<T*>(raw);}
#endif
return static_cast<T*>(::operator new(padded,std::align_val_t(align)));}void deallocate(T*p,std::size_t n)noexcept{if(!p)return;std::size_t padded=round_up(n*sizeof(T));std::size_t align=std::max<std::size_t>(alignof(T),Align);
#if CP_ALGO_USE_MMAP
if(padded>=MEGABYTE){munmap(p,padded);return;}
#endif
::operator delete(p,padded,std::align_val_t(align));}private:static constexpr std::size_t MEGABYTE=1<<20;static constexpr std::size_t round_up(std::size_t x)noexcept{return(x+Align-1)/Align*Align;}};template<typename T>using big_vector=std::vector<T,big_alloc<T>>;template<typename T>using big_basic_string=std::basic_string<T,std::char_traits<T>,big_alloc<T>>;template<typename T>using big_deque=std::deque<T,big_alloc<T>>;template<typename T>using big_stack=std::stack<T,big_deque<T>>;template<typename T>using big_queue=std::queue<T,big_deque<T>>;template<typename T>using big_priority_queue=std::priority_queue<T,big_vector<T>>;template<typename T>using big_forward_list=std::forward_list<T,big_alloc<T>>;using big_string=big_basic_string<char>;template<typename Key,typename Value,typename Compare=std::less<Key>>using big_map=std::map<Key,Value,Compare,big_alloc<std::pair<const Key,Value>>>;template<typename T,typename Compare=std::less<T>>using big_multiset=std::multiset<T,Compare,big_alloc<T>>;template<typename T,typename Compare=std::less<T>>using big_set=std::set<T,Compare,big_alloc<T>>;template<typename Ref,typename V=void>using big_generator=std::generator<Ref,V,big_alloc<std::byte>>;}namespace std::ranges{template<typename Ref,typename V>elements_of(cp_algo::big_generator<Ref,V>&&)->elements_of<cp_algo::big_generator<Ref,V>&&,cp_algo::big_alloc<std::byte>>;}
#line 5 "cp-algo/structures/stack_union.hpp"
#include <iterator>
#include <ranges>
namespace cp_algo::structures{template<class datatype>struct stack_union{stack_union(int n=0):head(n),next(1),data(1){}void push(int v,datatype const&vdata){next.push_back(head[v]);head[v]=(int)std::size(next)-1;data.push_back(vdata);}template<typename... Args>void emplace(int v,Args&&... vdata){next.push_back(head[v]);head[v]=(int)std::size(next)-1;data.emplace_back(std::forward<Args>(vdata)...);}void reserve(int m){data.reserve(m);next.reserve(m);}size_t size()const{return std::size(head);}size_t nodes()const{return std::size(data);}template<typename Su>struct _iterator{using value_type=std::conditional_t<std::is_const_v<Su>,const datatype,datatype>;using difference_type=std::ptrdiff_t;Su*su=nullptr;int sv=0;value_type&operator*()const{return su->data[sv];}_iterator&operator++(){sv=su->next[sv];return*this;}_iterator operator++(int){auto tmp=*this;++*this;return tmp;}friend bool operator==(_iterator const&it,std::default_sentinel_t){return it.sv==0;}};using iterator=_iterator<stack_union<datatype>>;using const_iterator=_iterator<const stack_union<datatype>>;auto operator[](this auto&&self,int v){using Iter=_iterator<std::remove_reference_t<decltype(self)>>;return std::ranges::subrange(Iter{&self,self.head[v]},std::default_sentinel);}big_vector<int>head,next;big_vector<datatype>data;};}
#line 7 "cp-algo/graph/base.hpp"
namespace cp_algo::graph{using edge_index=int;template<edge_type _edge_t=edge_base,graph_mode _mode=undirected>struct graph{using edge_t=_edge_t;static constexpr auto mode=_mode;using incidence_list=structures::stack_union<edge_index>;graph(int n,int v0=0):v0(v0),adj(n){}graph transpose()const{static_assert(mode==directed,"transpose is only defined for directed graphs");graph<edge_t,mode>gt(n(),v0);for(auto v:nodes()){for(auto e:outgoing(v)){gt.add_edge(edge(e).traverse(v),edge(e));}}return gt;}edge_index add_edge(node_index u,edge_t e){edge_index idx=(edge_index)size(E);E.push_back(e);adj.push(u,idx);if constexpr(mode==undirected){adj.push(e.traverse(u),idx);}return idx;}edge_index add_edge(node_index u,auto... Args){return add_edge(u,edge_t(u,Args...));}void read_edges(node_index m){adj.reserve(mode==undirected?2*m:m);for(edge_index i=0;i<m;i++){auto[u,e]=edge_t::read(v0);add_edge(u,e);}}auto outgoing(node_index v)const{return adj[v];}auto edges()const{return E|std::views::all;}auto nodes()const{return std::views::iota(node_index(0),n());}auto edge_indices()const{return std::views::iota(edge_index(0),m());}auto&&incidence_lists(this auto&&self){return self.adj;}auto&&edge(this auto&&self,edge_index e){return self.E[e];}node_index n()const{return(node_index)incidence_lists().size();}edge_index m()const{return(edge_index)edges().size();}private:node_index v0;big_vector<edge_t>E;incidence_list adj;};template<edge_type edge_t=edge_base>using digraph=graph<edge_t,directed>;template<weighted_edge_type edge_t=weighted_edge,graph_mode mode=undirected>using weighted_graph=graph<edge_t,mode>;template<weighted_edge_type edge_t=weighted_edge>using weighted_digraph=digraph<edge_t>;}
#line 1 "cp-algo/structures/dsu.hpp"
#line 4 "cp-algo/structures/dsu.hpp"
#include <numeric>
#line 6 "cp-algo/structures/dsu.hpp"
namespace cp_algo::structures{struct disjoint_set_union{disjoint_set_union(int n):par(n){std::iota(begin(par),end(par),0);}int get(int v){return v==par[v]?v:par[v]=get(par[v]);}bool uni(int a,int b){a=get(a);b=get(b);par[a]=b;return a!=b;}private:big_vector<int>par;};using dsu=disjoint_set_union;}
#line 1 "cp-algo/util/sort.hpp"
#line 1 "cp-algo/util/bit.hpp"
#line 1 "cp-algo/util/simd.hpp"
#include <experimental/simd>
#line 6 "cp-algo/util/simd.hpp"
#include <memory>
#if defined(__x86_64__) && !defined(CP_ALGO_DISABLE_AVX2)
#define CP_ALGO_SIMD_AVX2_TARGET _Pragma("GCC target(\"avx2\")")
#else
#define CP_ALGO_SIMD_AVX2_TARGET
#endif
#define CP_ALGO_SIMD_PRAGMA_PUSH \
_Pragma("GCC push_options")\CP_ALGO_SIMD_AVX2_TARGETCP_ALGO_SIMD_PRAGMA_PUSHnamespace cp_algo{template<typename T,size_t len>using simd[[gnu::vector_size(len*sizeof(T))]]=T;using i64x4=simd<int64_t,4>;using u64x4=simd<uint64_t,4>;using u32x8=simd<uint32_t,8>;using i32x4=simd<int32_t,4>;using u32x4=simd<uint32_t,4>;using i16x4=simd<int16_t,4>;using u8x32=simd<uint8_t,32>;using dx4=simd<double,4>;dx4 abs(dx4 a){return dx4{std::abs(a[0]),std::abs(a[1]),std::abs(a[2]),std::abs(a[3])};}static constexpr dx4 magic=dx4()+(3ULL<<51);inline i64x4 lround(dx4 x){return i64x4(x+magic)-i64x4(magic);}inline dx4 to_double(i64x4 x){return dx4(x+i64x4(magic))-magic;}inline dx4 round(dx4 a){return dx4{std::nearbyint(a[0]),std::nearbyint(a[1]),std::nearbyint(a[2]),std::nearbyint(a[3])};}inline u64x4 low32(u64x4 x){return x&uint32_t(-1);}inline auto swap_bytes(auto x){return decltype(x)(__builtin_shufflevector(u32x8(x),u32x8(x),1,0,3,2,5,4,7,6));}inline u64x4 montgomery_reduce(u64x4 x,uint32_t mod,uint32_t imod){
#ifdef __AVX2__
auto x_ninv=u64x4(_mm256_mul_epu32(__m256i(x),__m256i()+imod));x+=u64x4(_mm256_mul_epu32(__m256i(x_ninv),__m256i()+mod));
#else
auto x_ninv=u64x4(u32x8(low32(x))*imod);x+=x_ninv*uint64_t(mod);
#endif
return swap_bytes(x);}inline u64x4 montgomery_mul(u64x4 x,u64x4 y,uint32_t mod,uint32_t imod){
#ifdef __AVX2__
return montgomery_reduce(u64x4(_mm256_mul_epu32(__m256i(x),__m256i(y))),mod,imod);
#else
return montgomery_reduce(x*y,mod,imod);
#endif
}inline u32x8 montgomery_mul(u32x8 x,u32x8 y,uint32_t mod,uint32_t imod){return u32x8(montgomery_mul(u64x4(x),u64x4(y),mod,imod))|u32x8(swap_bytes(montgomery_mul(u64x4(swap_bytes(x)),u64x4(swap_bytes(y)),mod,imod)));}inline dx4 rotate_right(dx4 x){static constexpr u64x4 shuffler={3,0,1,2};return __builtin_shuffle(x,shuffler);}template<std::size_t Align=32>inline bool is_aligned(const auto*p)noexcept{return(reinterpret_cast<std::uintptr_t>(p)%Align)==0;}template<class Target>inline Target&vector_cast(auto&&p){return*reinterpret_cast<Target*>(std::assume_aligned<alignof(Target)>(&p));}}
#pragma GCC pop_options
#line 6 "cp-algo/util/bit.hpp"
#include <array>
#include <bit>
#if defined(__x86_64__) && !defined(CP_ALGO_DISABLE_AVX2)
#define CP_ALGO_BIT_OPS_TARGET _Pragma("GCC target(\"avx2,bmi,bmi2,lzcnt,popcnt\")")
#else
#define CP_ALGO_BIT_OPS_TARGET _Pragma("GCC target(\"bmi,bmi2,lzcnt,popcnt\")")
#endif
#define CP_ALGO_BIT_PRAGMA_PUSH \
_Pragma("GCC push_options")\CP_ALGO_BIT_OPS_TARGETCP_ALGO_BIT_PRAGMA_PUSHnamespace cp_algo{template<typename Uint>constexpr size_t bit_width=sizeof(Uint)*8;uint64_t mask(size_t n){return(1ULL<<n)-1;}size_t order_of_bit(auto x,size_t k){return k?std::popcount(x<<(bit_width<decltype(x)>-k)):0;}inline size_t kth_set_bit(uint64_t x,size_t k){return std::countr_zero(_pdep_u64(1ULL<<k,x));}template<int fl=0>void with_bit_floor(size_t n,auto&&callback){if constexpr(fl>=63){return;}else if(n>>(fl+1)){with_bit_floor<fl+1>(n,callback);}else{callback.template operator()<1ULL<<fl>();}}void with_bit_ceil(size_t n,auto&&callback){with_bit_floor(n,[&]<size_t N>(){if(N==n){callback.template operator()<N>();}else{callback.template operator()<N<<1>();}});}inline uint32_t read_bits(char const*p){return _mm256_movemask_epi8(__m256i(vector_cast<u8x32 const>(p[0])+(127-'0')));}inline uint64_t read_bits64(char const*p){return read_bits(p)|(uint64_t(read_bits(p+32))<<32);}inline void write_bits(char*p,uint32_t bits){static constexpr u8x32 shuffler={0,0,0,0,0,0,0,0,1,1,1,1,1,1,1,1,2,2,2,2,2,2,2,2,3,3,3,3,3,3,3,3};auto shuffled=u8x32(_mm256_shuffle_epi8(__m256i()+bits,__m256i(shuffler)));static constexpr u8x32 mask={1,2,4,8,16,32,64,128,1,2,4,8,16,32,64,128,1,2,4,8,16,32,64,128,1,2,4,8,16,32,64,128};for(int z=0;z<32;z++){p[z]=shuffled[z]&mask[z]?'1':'0';}}inline void write_bits64(char*p,uint64_t bits){write_bits(p,uint32_t(bits));write_bits(p+32,uint32_t(bits>>32));}}
#pragma GCC pop_options
#line 4 "cp-algo/util/sort.hpp"
#include <algorithm>
#line 8 "cp-algo/util/sort.hpp"
namespace cp_algo{template<size_t maxc>void count_sort(auto&a,auto&&proj){std::array<int,maxc>cnt={};for(auto&x:a){cnt[proj(x)]++;}std::partial_sum(begin(cnt),end(cnt),begin(cnt));auto res=a;for(auto const&it:a|std::views::reverse){res[--cnt[proj(it)]]=it;}a=std::move(res);}template<size_t maxc>void count_sort(auto&a){count_sort<maxc>(a,std::identity{});}void radix_sort(auto&a,auto&&proj){if(empty(a)){return;}auto[mn,mx]=std::ranges::minmax(a,{},proj);with_bit_floor<1>(size(a),[&]<size_t floor>(){constexpr int base=std::min<size_t>(floor,1<<16);for(int64_t i=1;i<=std::invoke(proj,mx)-std::invoke(proj,mn);i*=base){count_sort<base>(a,[&](auto const&x){return(std::invoke(proj,x)-std::invoke(proj,mn))/i%base;});}});}void radix_sort(auto&a){radix_sort(a,std::identity{});}}
#line 7 "cp-algo/graph/mst.hpp"
namespace cp_algo::graph{template<weighted_undirected_graph_type graph>std::pair<int64_t,big_vector<edge_index>>mst(graph const&g){struct edge{edge_index idx;node_index v;};big_vector<edge>edges;for(auto v:g.nodes()){for(auto e:g.outgoing(v)){if(v<g.edge(e).traverse(v)){edges.push_back({e,v});}}}radix_sort(edges,[&](auto e){return g.edge(e.idx).w;});structures::dsu me(g.n());int64_t total=0;big_vector<edge_index>mst;for(auto[idx,v]:edges){if(me.uni(v,g.edge(idx).traverse(v))){total+=g.edge(idx).w;mst.push_back(idx);}}return{total,mst};}}

Back to top page