library/graph/matching/WeightedBipartiteMatching.hpp

• View this file on GitHub
• Last update: 2025-11-24 18:49:07+09:00
• Include: #include "library/graph/matching/WeightedBipartiteMatching.hpp"

Depends on

• build/pch/stdc++.hpp
• library/flow/MCF.hpp
• library/graph/WeightedGraph.hpp

Verified with

• test/library-checker/Graph/AssignmentProblem.test.cpp

Code

#pragma once
#include "library/flow/MCF.hpp"
// 重みの最大化
template <typename TC> class WeightedBipartiteMatching {
    int A, B; // 左右の頂点数
    int S, T;
    MCF<int, TC> fl;

  public:
    WeightedBipartiteMatching(int A, int B)
        : A(A), B(B), S(A + B), T(A + B + 1), fl(A + B + 2, S, T) {
        for (int i = 0; i < A; i++)
            fl.add_arc(S, i, 1, 0);
        for (int j = 0; j < B; j++)
            fl.add_arc(A + j, T, 1, 0);
    }
    void add_edge(int u, int v, TC weight) {
        assert(0 <= u and u < A);
        assert(0 <= v and v < B);
        fl.add_arc(u, A + v, 1, -weight);
    }

    // first は重みの総和
    // second はマッチした各 [u,v,weight]
    std::pair<TC, std::vector<std::tuple<int, int, TC>>> solve() {
        auto [sum, ok] = fl.flow(std::min(A, B));
        std::vector<std::tuple<int, int, TC>> res;
        auto all_edge = fl.all_edge();
        for (int i = A + B; i < all_edge.size(); i++) {
            const auto &[from, to, flow, cost] = all_edge[i];
            if (flow)
                res.emplace_back(from, to - A, -cost);
        }
        return std::make_pair(-sum, res);
    }
};

#line 2 "library/graph/WeightedGraph.hpp"
template <typename T> struct WeightedEdge {
    WeightedEdge() = default;
    WeightedEdge(int from, int to, T weight)
        : from(from), to(to), weight(weight) {}
    int from, to;
    T weight;
    operator int() const { return to; }
};

template <typename T> struct WeightedGraph {
    int n;
    using weight_type = T;
    using edge_type = WeightedEdge<T>;
    std::vector<edge_type> edges;

  protected:
    std::vector<int> in_deg;
    bool prepared;
    class OutgoingEdges {
        WeightedGraph *g;
        int l, r;

      public:
        OutgoingEdges(WeightedGraph *g, int l, int r) : g(g), l(l), r(r) {}
        edge_type *begin() { return &(g->edges[l]); }
        edge_type *end() { return &(g->edges[r]); }
        edge_type &operator[](int i) { return g->edges[l + i]; }
        int size() const { return r - l; }
    };
    class ConstOutgoingEdges {
        const WeightedGraph *g;
        int l, r;

      public:
        ConstOutgoingEdges(const WeightedGraph *g, int l, int r)
            : g(g), l(l), r(r) {}
        const edge_type *begin() const { return &(g->edges[l]); }
        const edge_type *end() const { return &(g->edges[r]); }
        const edge_type &operator[](int i) const { return g->edges[l + i]; }
        int size() const { return r - l; }
    };

  public:
    OutgoingEdges operator[](int v) {
        assert(prepared);
        return {this, in_deg[v], in_deg[v + 1]};
    }
    const ConstOutgoingEdges operator[](int v) const {
        assert(prepared);
        return {this, in_deg[v], in_deg[v + 1]};
    }

    bool is_prepared() const { return prepared; }

    WeightedGraph() : n(0), in_deg(1, 0), prepared(false) {}
    WeightedGraph(int n) : n(n), in_deg(n + 1, 0), prepared(false) {}
    WeightedGraph(int n, int m, bool directed = false, int indexed = 1)
        : n(n), in_deg(n + 1, 0), prepared(false) {
        scan(m, directed, indexed);
    }

    void resize(int n) { n = n; }

    void add_arc(int from, int to, T weight) {
        assert(!prepared);
        assert(0 <= from and from < n and 0 <= to and to < n);
        edges.emplace_back(from, to, weight);
        in_deg[from + 1]++;
    }
    void add_edge(int u, int v, T weight) {
        add_arc(u, v, weight);
        add_arc(v, u, weight);
    }
    void add_arc(const edge_type &e) { add_arc(e.from, e.to, e.weight); }
    void add_edge(const edge_type &e) { add_edge(e.from, e.to, e.weight); }

    void scan(int m, bool directed = false, int indexed = 1) {
        edges.reserve(directed ? m : 2 * m);
        while (m--) {
            int u, v;
            std::cin >> u >> v;
            u -= indexed;
            v -= indexed;
            T weight;
            std::cin >> weight;
            if (directed)
                add_arc(u, v, weight);
            else
                add_edge(u, v, weight);
        }
        build();
    }

    void build() {
        assert(!prepared);
        prepared = true;
        for (int v = 0; v < n; v++)
            in_deg[v + 1] += in_deg[v];
        std::vector<edge_type> new_edges(in_deg.back());
        auto counter = in_deg;
        for (auto &&e : edges)
            new_edges[counter[e.from]++] = e;
        edges = new_edges;
    }

    void graph_debug() const {
#ifndef __DEBUG
        return;
#endif
        assert(prepared);
        for (int from = 0; from < n; from++) {
            std::cerr << from << ";";
            for (int i = in_deg[from]; i < in_deg[from + 1]; i++)
                std::cerr << "(" << edges[i].to << "," << edges[i].weight
                          << ")";
            std::cerr << "\n";
        }
    }
};
#line 3 "library/flow/MCF.hpp"
#define REP_(i, n) for (int i = 0; i < (n); i++)
template <typename TF, typename TC> class MCF {
    struct EdgeInfo {
        TF cap;
        TC cost;
        int rev;
    };
    int n, s, t;
    WeightedGraph<EdgeInfo> G;
    std::vector<TC> potential, dist;
    static constexpr TC INF = std::numeric_limits<TC>::max() / 2;
    std::vector<std::pair<int, int>> pre,
        edge_memo; // pre[v]=[u,i] : G[u][i] で v に来た
    std::vector<int> in_deg, out_deg;

    bool negative, dag;

    template <typename T> bool chmin(T &a, const T &b) {
        return (a > b and (a = b, true));
    }
    bool SP_update(int from, int edge_id) {
        if (dist[from] == INF)
            return false;
        const auto &e = G[from][edge_id];
        if ((e.weight).cap == 0)
            return false;
        if (chmin(dist[e.to], dist[from] + (e.weight).cost + potential[from] -
                                  potential[e.to])) {
            pre[e.to] = {from, edge_id};
            return true;
        }
        return false;
    }

    std::priority_queue<std::pair<TC, int>, std::vector<std::pair<TC, int>>,
                        std::greater<std::pair<TC, int>>>
        que;
    void dijkstra() { // dist[i]:sから残余グラフで辺の重みによるiへの最短路
                      // となるようにdistを作る
        std::ranges::fill(dist, INF);
        dist[s] = 0;
        que.emplace(0, s);
        while (que.size()) {
            const auto [now, v] = que.top();
            que.pop();
            if (dist[v] < now)
                continue;
            REP_(i, G[v].size())
            if (SP_update(v, i))
                que.emplace(dist[G[v][i].to], G[v][i].to);
        }
    }

    void DAG() {
        negative = false;
        std::ranges::fill(dist, INF);
        dist[s] = 0;
        std::queue<int> que;
        REP_(i, n) if (!in_deg[i]) que.push(i);
        while (que.size()) {
            int v = que.front();
            que.pop();
            REP_(i, G[v].size()) {
                SP_update(v, i);
                if (!--in_deg[G[v][i].to])
                    que.push(G[v][i].to);
            }
        }
    }

    void BellmanFord() {
        negative = false;
        std::ranges::fill(dist, INF);
        dist[s] = 0;
        REP_(_, n) {
            bool update = false;
            REP_(v, n)
            if (dist[v] < INF)
                REP_(i, G[v].size()) if (SP_update(v, i)) update = true;
            if (!update)
                return;
        }
        assert(false); // 負閉路
    }

  public:
    MCF() {}
    MCF(int n_, int s_ = 0, int t_ = -1)
        : n(n_), G(n_), potential(n_, 0), dist(n_), pre(n_), in_deg(n_, 0),
          out_deg(n_, 0), negative(false), dag(true), s(s_), t(t_) {
        if (t < 0)
            t = n - 1;
    }
    void use_bellman_ford() { dag = false; }

    TF operator[](const int edge_id) const {
        assert(G.is_prepared());
        const auto &[from, id] = edge_memo[edge_id];
        return G.edge[from][id].weight.cap;
    }
    std::vector<std::tuple<int, int, TF, TC>> all_edge() {
        assert(G.is_prepared());
        std::vector<std::tuple<int, int, TF, TC>> res;
        res.reserve(edge_memo.size());
        for (const auto &[v, id] : edge_memo) {
            const auto &[to, from, weight] = G[v][id];
            res.emplace_back(from, to, weight.cap, -weight.cost);
        }
        return res;
    }

    void add_arc(int from, int to, TF cap, TC cost) {
        G.add_arc(from, to, {cap, cost, out_deg[to]});
        G.add_arc(to, from, {0, -cost, out_deg[from]++});
        edge_memo.emplace_back(to, out_deg[to]++);
        if (cap > 0) {
            in_deg[to]++;
            negative |= cost < 0;
        }
    }

    std::pair<TC, bool>
    flow(TF f = std::numeric_limits<
             TF>::max()) { // second が 0
                           // で返ってきた場合はそもそも最大流がfに達しない
        if (!G.is_prepared())
            G.build();
        TC res = 0;
        std::ranges::fill(
            potential,
            0); // 一番最初は負のコストの辺が無いからポテンシャルは0にしていい
        while (f > 0) {
            if (negative)
                if (dag)
                    DAG();
                else
                    BellmanFord();
            else
                dijkstra();
            if (dist[t] == INF)
                return std::make_pair(res, false);
            REP_(v, n) if (dist[v] < INF) potential[v] += dist[v];
            TF d = f; // d:今回流す量
            for (int v = t; v != s; v = pre[v].first)
                chmin(d, (G[pre[v].first][pre[v].second].weight).cap);
            f -= d;
            res += potential[t] * d;
            for (int v = t; v != s; v = pre[v].first) {
                auto &[cap, cost, rev] = G[pre[v].first][pre[v].second].weight;
                cap -= d;
                (G[v][rev].weight).cap += d;
            }
        } // このループを抜けてるならf流れてる
        return std::make_pair(res, true);
    }

    std::pair<TC, bool> st_flow(int s_, int t_,
                                TF lim = std::numeric_limits<TF>::max() / 2) {
        s = s_;
        t = t_;
        return flow(lim);
    }
};
#undef REP_
#line 3 "library/graph/matching/WeightedBipartiteMatching.hpp"
// 重みの最大化
template <typename TC> class WeightedBipartiteMatching {
    int A, B; // 左右の頂点数
    int S, T;
    MCF<int, TC> fl;

  public:
    WeightedBipartiteMatching(int A, int B)
        : A(A), B(B), S(A + B), T(A + B + 1), fl(A + B + 2, S, T) {
        for (int i = 0; i < A; i++)
            fl.add_arc(S, i, 1, 0);
        for (int j = 0; j < B; j++)
            fl.add_arc(A + j, T, 1, 0);
    }
    void add_edge(int u, int v, TC weight) {
        assert(0 <= u and u < A);
        assert(0 <= v and v < B);
        fl.add_arc(u, A + v, 1, -weight);
    }

    // first は重みの総和
    // second はマッチした各 [u,v,weight]
    std::pair<TC, std::vector<std::tuple<int, int, TC>>> solve() {
        auto [sum, ok] = fl.flow(std::min(A, B));
        std::vector<std::tuple<int, int, TC>> res;
        auto all_edge = fl.all_edge();
        for (int i = A + B; i < all_edge.size(); i++) {
            const auto &[from, to, flow, cost] = all_edge[i];
            if (flow)
                res.emplace_back(from, to - A, -cost);
        }
        return std::make_pair(-sum, res);
    }
};

Back to top page