graph/binary_optimization.cpp

View this file on GitHub
Last update: 2022-03-31 10:32:23+09:00

Depends on

Code

#pragma once
#include <cassert>
#include <limits>
#include <vector>
#include "../utility/index_offset.cpp"
#include "../utility/rep.cpp"
#include "dinic.cpp"

template <class T> class BinaryOptimization {
    Dinic<T> graph;
    T constant;
    std::vector<int> var_id;
    int src, dst;

  public:
    BinaryOptimization() : graph(), constant(), var_id(), src(graph.add_vertex()), dst(graph.add_vertex()) {}
    explicit BinaryOptimization(const int n) : BinaryOptimization() { add_variables(n); }

    int count_variables() const { return var_id.size(); }

    int add_variable() {
        var_id.push_back(graph.add_vertex());
        return count_variables() - 1;
    }
    IndexOffset add_variables(int n) {
        IndexOffset ret(count_variables(), n);
        while (n--) var_id.push_back(graph.add_vertex());
        return ret;
    }

    void add_cost(const T& c) { constant += c; }

    template <class F> void add_cost(const int x, const F& f) {
        assert(0 <= x and x < count_variables());
        const T a = f(0), b = f(1);
        if (a <= b) {
            graph.add_edge(src, var_id[x], b - a);
            add_cost(a);
        } else {
            graph.add_edge(var_id[x], dst, a - b);
            add_cost(b);
        }
    }

    template <class F> void add_cost(const int x, const int y, const F& f) {
        assert(0 <= x and x < count_variables());
        assert(0 <= y and y < count_variables());
        const T a = f(0, 0), b = f(0, 1), c = f(1, 0), d = f(1, 1);
        assert(b + c >= a + d);
        add_cost(a);
        add_cost(x, [&](const bool v) { return v ? c - a : 0; });
        add_cost(y, [&](const bool v) { return v ? d - c : 0; });
        graph.add_edge(var_id[x], var_id[y], (b + c) - (a + d));
    }

    void disable(const int x, const bool f) {
        assert(0 <= x and x < count_variables());
        if (f)
            graph.add_edge(src, var_id[x], std::numeric_limits<T>::max());
        else
            graph.add_edge(var_id[x], dst, std::numeric_limits<T>::max());
    }

    void disable(const int x, const int y, const bool f, const bool g) {
        assert(0 <= x and x < count_variables());
        assert(0 <= y and y < count_variables());
        assert(f ^ g);
        if (f)
            graph.add_edge(var_id[y], var_id[x], std::numeric_limits<T>::max());
        else
            graph.add_edge(var_id[x], var_id[y], std::numeric_limits<T>::max());
    }

    T solve() { return constant + graph.flow(src, dst); }
    std::vector<char> restore() {
        const std::vector<char> cut = graph.min_cut(src);
        const int n = count_variables();
        std::vector<char> ret(n);
        for (const int i : rep(n)) ret[i] = !cut[var_id[i]];
        return ret;
    }
};

#line 2 "graph/binary_optimization.cpp"
#include <cassert>
#include <limits>
#include <vector>
#line 3 "utility/index_offset.cpp"

class IndexOffset {
    int offset, len;

  public:
    constexpr IndexOffset() noexcept : offset(), len() {}
    explicit constexpr IndexOffset(const int o, const int l) noexcept : offset(o), len(l) {}
    constexpr int size() const { return len; }
    constexpr int operator[](const int i) const noexcept {
        assert(i < len);
        return offset + i;
    }
    constexpr int to_idx(const int i) const noexcept {
        assert(offset <= i and i < offset + len);
        return i - offset;
    }
};
#line 2 "utility/rep.cpp"
#include <algorithm>

class Range {
    struct Iter {
        int itr;
        constexpr Iter(const int pos) noexcept : itr(pos) {}
        constexpr void operator++() noexcept { ++itr; }
        constexpr bool operator!=(const Iter& other) const noexcept { return itr != other.itr; }
        constexpr int operator*() const noexcept { return itr; }
    };
    const Iter first, last;

  public:
    explicit constexpr Range(const int first, const int last) noexcept : first(first), last(std::max(first, last)) {}
    constexpr Iter begin() const noexcept { return first; }
    constexpr Iter end() const noexcept { return last; }
};

constexpr Range rep(const int l, const int r) noexcept { return Range(l, r); }
constexpr Range rep(const int n) noexcept { return Range(0, n); }
#line 5 "graph/dinic.cpp"
#include <type_traits>
#line 3 "internal/simple_queue.cpp"

namespace proconlib {

template <class T> class SimpleQueue {
    std::vector<T> payload;
    int pos;

  public:
    SimpleQueue() : payload(), pos(0) {}
    explicit SimpleQueue(const int n) : SimpleQueue() { reserve(n); }

    int size() const { return (int)payload.size() - pos; }
    bool empty() const { return size() == 0; }
    T& front() { return payload[pos]; }

    void push(const T& t) { payload.push_back(t); }
    void push(T&& t) { payload.push_back(std::move(t)); }
    void pop() { pos++; }

    void reserve(int n) { payload.reserve(n); }
    void clear() {
        payload.clear();
        pos = 0;
    }
};

}  // namespace proconlib
#line 3 "utility/rec_lambda.cpp"
#include <utility>

template <class F> struct RecursiveLambda : private F {
    explicit constexpr RecursiveLambda(F&& f) : F(std::forward<F>(f)) {}
    template <class... Args> constexpr decltype(auto) operator()(Args&&... args) const {
        return F::operator()(*this, std::forward<Args>(args)...);
    }
};

template <class F> constexpr decltype(auto) rec_lambda(F&& f) { return RecursiveLambda<F>(std::forward<F>(f)); }
#line 11 "graph/dinic.cpp"

template <class Flow, std::enable_if_t<std::is_integral_v<Flow>>* = nullptr> class Dinic {
  public:
    struct Edge {
        int src, dst;
        Flow flow, cap;
    };

  private:
    int node_count;
    std::vector<Edge> graph;

  public:
    Dinic() : node_count(0), graph() {}
    explicit Dinic(const int n) : node_count(n), graph() {}

    int size() const { return node_count; }
    int edge_count() const { return graph.size(); }

    int add_vertex() { return node_count++; }
    IndexOffset add_vertices(const int n) {
        assert(n >= 0);
        const IndexOffset ret(size(), n);
        node_count += n;
        return ret;
    }

    const Edge& get_edge(const int i) const {
        assert(0 <= i and i < edge_count());
        return graph[i];
    }
    int add_edge(const int src, const int dst, const Flow& cap) {
        assert(0 <= src and src < size());
        assert(0 <= dst and dst < size());
        assert(cap >= 0);
        graph.push_back(Edge{src, dst, 0, cap});
        return edge_count() - 1;
    }

    Flow flow(const int src, const int dst) { return flow(src, dst, std::numeric_limits<Flow>::max()); }
    Flow flow(const int src, const int dst, const Flow& flow_limit) {
        assert(0 <= src and src < size());
        assert(0 <= dst and dst < size());
        assert(src != dst);
        const int n = size();
        const int m = edge_count();
        struct E {
            int dst, rev;
            Flow cap;
        };
        std::vector<E> edge(2 * m);
        std::vector<int> start(n + 1), eidx(m);
        {
            std::vector<int> deg(n), reidx(m);
            for (const int i : rep(m)) {
                eidx[i] = deg[graph[i].src]++;
                reidx[i] = deg[graph[i].dst]++;
            }
            for (const int i : rep(n)) start[i + 1] = start[i] + deg[i];
            for (const int i : rep(m)) {
                const auto& e = graph[i];
                const int u = e.src, v = e.dst;
                eidx[i] += start[u];
                reidx[i] += start[v];
                edge[eidx[i]] = {v, reidx[i], e.cap - e.flow};
                edge[reidx[i]] = {u, eidx[i], e.flow};
            }
        }
        std::vector<int> level(n), iter(n);
        proconlib::SimpleQueue<int> que;
        const auto bfs = [&] {
            std::fill(level.begin(), level.end(), n);
            level[src] = 0;
            while (!que.empty()) que.pop();
            que.push(src);
            while (!que.empty()) {
                const int u = que.front();
                que.pop();
                for (const int i : rep(start[u], start[u + 1])) {
                    const auto& e = edge[i];
                    if (e.cap == 0 or level[e.dst] < n) continue;
                    level[e.dst] = level[u] + 1;
                    if (e.dst == dst) return;
                    que.push(e.dst);
                }
            }
        };
        const auto dfs = rec_lambda([&](auto&& dfs, const int u, const Flow& ub) -> Flow {
            if (u == src) return ub;
            Flow ret = 0;
            for (int& i = iter[u]; i < start[u + 1]; i += 1) {
                auto& e = edge[i];
                auto& re = edge[e.rev];
                if (level[u] <= level[e.dst] or re.cap == 0) continue;
                const Flow d = dfs(e.dst, std::min(ub - ret, re.cap));
                if (d == 0) continue;
                e.cap += d;
                re.cap -= d;
                ret += d;
                if (ret == ub) return ret;
            }
            level[u] = n;
            return ret;
        });
        Flow ret = 0;
        while (ret < flow_limit) {
            bfs();
            if (level[dst] == n) break;
            std::copy(start.begin(), start.begin() + n, iter.begin());
            const Flow f = dfs(dst, flow_limit - ret);
            if (f == 0) break;
            ret += f;
        }
        for (const int i : rep(m)) graph[i].flow = graph[i].cap - edge[eidx[i]].cap;
        return ret;
    }

    std::vector<char> min_cut(const int src) const {
        assert(0 <= src and src < size());
        const int n = size();
        std::vector<std::vector<int>> adj(n);
        for (const auto& e : graph) {
            if (e.flow < e.cap) adj[e.src].push_back(e.dst);
            if (e.flow > 0) adj[e.dst].push_back(e.src);
        }
        std::vector<char> ret(n);
        proconlib::SimpleQueue<int> que;
        que.push(src);
        ret[src] = true;
        while (!que.empty()) {
            const int u = que.front();
            que.pop();
            for (const int v : adj[u]) {
                if (!ret[v]) {
                    ret[v] = true;
                    que.push(v);
                }
            }
        }
        return ret;
    }
};
#line 8 "graph/binary_optimization.cpp"

template <class T> class BinaryOptimization {
    Dinic<T> graph;
    T constant;
    std::vector<int> var_id;
    int src, dst;

  public:
    BinaryOptimization() : graph(), constant(), var_id(), src(graph.add_vertex()), dst(graph.add_vertex()) {}
    explicit BinaryOptimization(const int n) : BinaryOptimization() { add_variables(n); }

    int count_variables() const { return var_id.size(); }

    int add_variable() {
        var_id.push_back(graph.add_vertex());
        return count_variables() - 1;
    }
    IndexOffset add_variables(int n) {
        IndexOffset ret(count_variables(), n);
        while (n--) var_id.push_back(graph.add_vertex());
        return ret;
    }

    void add_cost(const T& c) { constant += c; }

    template <class F> void add_cost(const int x, const F& f) {
        assert(0 <= x and x < count_variables());
        const T a = f(0), b = f(1);
        if (a <= b) {
            graph.add_edge(src, var_id[x], b - a);
            add_cost(a);
        } else {
            graph.add_edge(var_id[x], dst, a - b);
            add_cost(b);
        }
    }

    template <class F> void add_cost(const int x, const int y, const F& f) {
        assert(0 <= x and x < count_variables());
        assert(0 <= y and y < count_variables());
        const T a = f(0, 0), b = f(0, 1), c = f(1, 0), d = f(1, 1);
        assert(b + c >= a + d);
        add_cost(a);
        add_cost(x, [&](const bool v) { return v ? c - a : 0; });
        add_cost(y, [&](const bool v) { return v ? d - c : 0; });
        graph.add_edge(var_id[x], var_id[y], (b + c) - (a + d));
    }

    void disable(const int x, const bool f) {
        assert(0 <= x and x < count_variables());
        if (f)
            graph.add_edge(src, var_id[x], std::numeric_limits<T>::max());
        else
            graph.add_edge(var_id[x], dst, std::numeric_limits<T>::max());
    }

    void disable(const int x, const int y, const bool f, const bool g) {
        assert(0 <= x and x < count_variables());
        assert(0 <= y and y < count_variables());
        assert(f ^ g);
        if (f)
            graph.add_edge(var_id[y], var_id[x], std::numeric_limits<T>::max());
        else
            graph.add_edge(var_id[x], var_id[y], std::numeric_limits<T>::max());
    }

    T solve() { return constant + graph.flow(src, dst); }
    std::vector<char> restore() {
        const std::vector<char> cut = graph.min_cut(src);
        const int n = count_variables();
        std::vector<char> ret(n);
        for (const int i : rep(n)) ret[i] = !cut[var_id[i]];
        return ret;
    }
};