m1une's library

This documentation is automatically generated by online-judge-tools/verification-helper

View on GitHub

:heavy_check_mark: K-Value Project Selection
(optimization/k_project_selection.hpp)

Overview

KProjectSelection<T> assigns one ordered value to every project and maximizes the sum of unary gains and graph-cut-representable pairwise gains. Project i has values 0, 1, ..., value_count(i) - 1; different projects may have different numbers of values.

This is the ordered, k-valued generalization of binary project selection. It is useful for labeling and discrete optimization problems whose pairwise terms prefer coordinated increases. The model is reduced to one minimum s-t cut.

Objective Terms

All calls are additive: calling add_gain several times adds every supplied term to the same objective.

Method Contribution
void add_gain(int i, const std::vector<T>& gains) Adds gains[x] when project i takes value x. The vector must have value_count(i) entries.
void add_gain(int i, int j, const std::vector<std::vector<T>>& gains) Adds gains[x][y] when distinct projects i and j take values x and y. The matrix must have shape value_count(i) by value_count(j).

Domain of Gains

Every entry in a unary gain vector or pairwise gain matrix may be negative, zero, or positive, as long as it is representable by T. A negative entry represents a cost for the corresponding value or pair of values. Consequently, the optimal result.max_gain may also be negative.

Unary gain vectors have no further algebraic restriction. Pairwise matrix entries do not need to be non-negative, but their discrete mixed differences must be non-negative. Equivalently, a pairwise gain matrix g must be supermodular on the ordered values:

\[g(x,y) + g(x+1,y+1) \ge g(x+1,y) + g(x,y+1)\]

for every adjacent pair of rows and columns. The implementation checks this condition with assertions. Each mixed difference, every accumulated value, and the final answer must also fit the numeric requirements stated below.

Equivalently, if the problem is written using costs c = -g, every pairwise cost matrix must be Monge. An arbitrary pairwise table cannot generally be optimized by one minimum cut.

To minimize costs, negate every cost table before adding it, maximize the resulting gain, and negate result.max_gain.

Construction and Bounds

Method Signature / Meaning
Default constructor Creates a model with no projects.
Heterogeneous constructor explicit KProjectSelection(std::vector<int> value_counts) creates one project for every entry. Every count must be positive.
Uniform constructor KProjectSelection(int project_count, int value_count) gives every project the same positive number of values.
size int size() const returns the number of projects.
value_count int value_count(int i) const returns the number of values of project i.
force_value void force_value(int i, int x) requires value x.
force_value_at_least void force_value_at_least(int i, int lower_bound) requires the value to be at least the bound.
force_value_at_most void force_value_at_most(int i, int upper_bound) requires the value to be at most the bound.

Bounds are inclusive. Contradictory bounds make the model infeasible rather than silently choosing a value.

Result

KProjectSelectionResult<T> contains:

Member / Method Type / Signature Meaning
feasible bool Whether all forced-value constraints can be satisfied.
max_gain T Maximum total gain; meaningful only when feasible.
values std::vector<int> One optimal value for each project.
is_feasible bool is_feasible() const Returns feasible.

solve() has signature KProjectSelectionResult<T> solve() const. It does not mutate the model, and it may return any optimum when several assignments tie.

Reduction and Complexity

For every project value boundary d, the reduction creates a Boolean threshold vertex representing value >= d. Hard implications ensure that the selected thresholds form a prefix. Unary tables become gains on these thresholds.

For a pairwise table, each non-negative discrete mixed difference becomes a reward for selecting two threshold vertices. Supermodularity is exactly the condition that makes all such rewards graph-cut-representable.

Let

Method Time Extra stored memory
Heterogeneous or uniform constructor $O(N + S)$ $O(N + S)$
size, value_count $O(1)$ $O(1)$
Unary add_gain $O(K_i)$ $O(K_i)$
Pairwise add_gain $O(K_i K_j)$ $O(K_i K_j)$ in the worst case
Forced-value or bound method $O(1)$ $O(1)$
solve General-case $O((S+2)^2 M)$ $O(S + M)$

The solve bound follows the underlying MaxFlow<T> implementation. Threshold reconstruction takes an additional $O(S)$ time.

Numeric Requirements

T must be a signed integral type no wider than 64 bits; long long is recommended. Every adjacent gain difference, discrete mixed difference, sum of finite cut capacities, accumulated constant, and final answer must fit in T. The implementation uses a wider intermediate for checking individual table differences and assertions for the remaining range requirements.

Example

The following model has three values for project 0 and two for project 1. Its pairwise table is supermodular.

#include "optimization/k_project_selection.hpp"
#include <iostream>
#include <vector>

int main() {
    m1une::opt::KProjectSelection<long long> solver(
        std::vector<int>{3, 2}
    );
    solver.add_gain(0, std::vector<long long>{1, 4, 2});
    solver.add_gain(1, std::vector<long long>{-2, 3});

    std::vector<std::vector<long long>> together(
        3,
        std::vector<long long>(2)
    );
    together[0] = std::vector<long long>{0, 1};
    together[1] = std::vector<long long>{-1, 3};
    together[2] = std::vector<long long>{1, 7};
    solver.add_gain(0, 1, together);

    auto result = solver.solve();
    if (!result.is_feasible()) return 0;
    std::cout << result.max_gain << '\n';  // 12
    std::cout << result.values[0] << ' ' << result.values[1] << '\n';  // 2 1
}

Depends on

Required by

Verified with

Code

#ifndef M1UNE_OPTIMIZATION_K_PROJECT_SELECTION_HPP
#define M1UNE_OPTIMIZATION_K_PROJECT_SELECTION_HPP 1

#include <cassert>
#include <cstddef>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

#include "project_selection.hpp"

namespace m1une {
namespace opt {

template <class T>
struct KProjectSelectionResult {
    bool feasible;
    T max_gain;
    std::vector<int> values;

    bool is_feasible() const {
        return feasible;
    }
};

template <class T>
class KProjectSelection {
    static_assert(std::is_integral_v<T> && std::is_signed_v<T>);
    static_assert(sizeof(T) <= sizeof(long long));

    using Wide = __int128_t;

    std::vector<int> _value_counts;
    std::vector<int> _first_threshold;
    ProjectSelection<T> _binary;
    T _constant = T();

    static int threshold_count(const std::vector<int>& value_counts) {
        assert(value_counts.size() <=
               std::size_t(std::numeric_limits<int>::max()));
        long long count = 0;
        for (int value_count : value_counts) {
            assert(value_count >= 1);
            count += value_count - 1;
            assert(count <= std::numeric_limits<int>::max());
        }
        return int(count);
    }

    static std::vector<int> repeated_value_counts(
        int project_count,
        int value_count
    ) {
        assert(project_count >= 0);
        assert(value_count >= 1);
        return std::vector<int>(project_count, value_count);
    }

    void assert_project(int project) const {
        (void)project;
        assert(0 <= project && project < size());
    }

    int threshold(int project, int value) const {
        assert_project(project);
        (void)value;
        assert(1 <= value && value < _value_counts[project]);
        return _first_threshold[project] + value - 1;
    }

    static T narrow(Wide value) {
        assert(Wide(std::numeric_limits<T>::lowest()) <= value);
        assert(value <= Wide(std::numeric_limits<T>::max()));
        return T(value);
    }

    void add_constant(T gain) {
        _constant = narrow(Wide(_constant) + gain);
    }

    void add_threshold_gain(int project, int value, Wide gain) {
        if (gain == 0) return;
        _binary.add_gain(threshold(project, value), narrow(gain));
    }

   public:
    KProjectSelection() : KProjectSelection(std::vector<int>()) {}

    explicit KProjectSelection(std::vector<int> value_counts)
        : _value_counts(std::move(value_counts)),
          _first_threshold(_value_counts.size()),
          _binary(threshold_count(_value_counts)) {
        int first = 0;
        for (int project = 0; project < size(); project++) {
            _first_threshold[project] = first;
            first += _value_counts[project] - 1;
        }

        for (int project = 0; project < size(); project++) {
            for (int value = 2; value < _value_counts[project]; value++) {
                _binary.add_hard_implication(
                    threshold(project, value),
                    threshold(project, value - 1)
                );
            }
        }
    }

    KProjectSelection(int project_count, int value_count)
        : KProjectSelection(repeated_value_counts(project_count, value_count)) {}

    int size() const {
        return int(_value_counts.size());
    }

    int value_count(int project) const {
        assert_project(project);
        return _value_counts[project];
    }

    void add_gain(int project, const std::vector<T>& gains) {
        assert_project(project);
        assert(int(gains.size()) == _value_counts[project]);
        add_constant(gains[0]);
        for (int value = 1; value < _value_counts[project]; value++) {
            add_threshold_gain(
                project,
                value,
                Wide(gains[value]) - gains[value - 1]
            );
        }
    }

    void add_gain(
        int project_a,
        int project_b,
        const std::vector<std::vector<T>>& gains
    ) {
        assert_project(project_a);
        assert_project(project_b);
        assert(project_a != project_b);
        const int count_a = _value_counts[project_a];
        const int count_b = _value_counts[project_b];
        assert(int(gains.size()) == count_a);
        for (const auto& row : gains) assert(int(row.size()) == count_b);

        add_constant(gains[0][0]);
        for (int value_a = 1; value_a < count_a; value_a++) {
            add_threshold_gain(
                project_a,
                value_a,
                Wide(gains[value_a][0]) - gains[value_a - 1][0]
            );
        }
        for (int value_b = 1; value_b < count_b; value_b++) {
            add_threshold_gain(
                project_b,
                value_b,
                Wide(gains[0][value_b]) - gains[0][value_b - 1]
            );
        }

        for (int value_a = 1; value_a < count_a; value_a++) {
            for (int value_b = 1; value_b < count_b; value_b++) {
                Wide mixed =
                    Wide(gains[value_a][value_b])
                    - gains[value_a - 1][value_b]
                    - gains[value_a][value_b - 1]
                    + gains[value_a - 1][value_b - 1];
                assert(mixed >= 0);
                T gain = narrow(mixed);
                if (gain == T()) continue;
                int threshold_a = threshold(project_a, value_a);
                int threshold_b = threshold(project_b, value_b);
                _binary.add_gain(threshold_a, gain);
                _binary.add_penalty(threshold_a, threshold_b, gain);
            }
        }
    }

    void force_value(int project, int value) {
        assert_project(project);
        assert(0 <= value && value < _value_counts[project]);
        force_value_at_least(project, value);
        force_value_at_most(project, value);
    }

    void force_value_at_least(int project, int lower_bound) {
        assert_project(project);
        assert(0 <= lower_bound && lower_bound < _value_counts[project]);
        if (lower_bound > 0) {
            _binary.force_selected(threshold(project, lower_bound));
        }
    }

    void force_value_at_most(int project, int upper_bound) {
        assert_project(project);
        assert(0 <= upper_bound && upper_bound < _value_counts[project]);
        if (upper_bound + 1 < _value_counts[project]) {
            _binary.force_unselected(threshold(project, upper_bound + 1));
        }
    }

    KProjectSelectionResult<T> solve() const {
        auto binary_result = _binary.solve();
        KProjectSelectionResult<T> result;
        result.feasible = binary_result.feasible;
        result.max_gain = T();
        result.values.assign(size(), 0);
        if (!result.feasible) return result;

        result.max_gain = narrow(Wide(_constant) + binary_result.max_gain);
        for (int project = 0; project < size(); project++) {
            for (int value = 1; value < _value_counts[project]; value++) {
                if (!binary_result.selected[threshold(project, value)]) break;
                result.values[project] = value;
            }
        }
        return result;
    }
};

}  // namespace opt
}  // namespace m1une

#endif  // M1UNE_OPTIMIZATION_K_PROJECT_SELECTION_HPP
#line 1 "optimization/k_project_selection.hpp"



#include <cassert>
#include <cstddef>
#include <limits>
#include <type_traits>
#include <utility>
#include <vector>

#line 1 "optimization/project_selection.hpp"



#line 9 "optimization/project_selection.hpp"

#line 1 "graph/flow/max_flow.hpp"



#include <algorithm>
#line 9 "graph/flow/max_flow.hpp"

namespace m1une {
namespace flow {

template <class Cap>
struct MaxFlow {
    struct Edge {
        int from;
        int to;
        Cap cap;
        Cap flow;
    };

   private:
    struct InternalEdge {
        int to;
        int rev;
        Cap cap;
    };

    struct Position {
        int from;
        int edge;
    };

    int _n;
    std::vector<Position> _pos;
    std::vector<std::vector<InternalEdge>> _g;

    Cap highest_label_preflow_push(int s, int t) {
        const int dead = 2 * _n;
        const int unreachable = _n + 1;
        std::vector<Cap> excess(_n, Cap(0));
        std::vector<int> state(8 * std::size_t(_n) + 2);
        int* height = state.data();
        int* height_count = height + _n;
        int* current = height_count + dead + 1;
        int* queue = current + _n;
        int* next = queue + _n;
        int* bucket_head = next + _n;
        std::vector<char> active(_n, false);
        int highest = -1;
        long long work = 0;
        const long long arc_count =
            2LL * static_cast<long long>(_pos.size());
        const long long work_limit = std::max(1LL, 4 * arc_count + _n);

        auto activate = [&](int v) {
            if (v == s || v == t || active[v] || excess[v] == Cap(0) ||
                height[v] >= dead) {
                return;
            }
            active[v] = true;
            next[v] = bucket_head[height[v]];
            bucket_head[height[v]] = v;
            highest = std::max(highest, height[v]);
        };

        auto rebuild_buckets = [&]() {
            std::fill(bucket_head, bucket_head + dead + 1, -1);
            std::fill(active.begin(), active.end(), false);
            highest = -1;
            for (int v = 0; v < _n; v++) activate(v);
        };

        auto global_relabel = [&]() {
            std::fill(height, height + _n, unreachable);
            std::fill(height_count, height_count + dead + 1, 0);
            std::fill(current, current + _n, 0);
            int head = 0;
            int tail = 0;
            height[t] = 0;
            height[s] = _n;
            queue[tail++] = t;
            while (head != tail) {
                int v = queue[head++];
                for (const auto& e : _g[v]) {
                    if (e.to == s || height[e.to] != unreachable) continue;
                    const auto& reverse = _g[e.to][e.rev];
                    if (reverse.cap == Cap(0)) continue;
                    height[e.to] = height[v] + 1;
                    queue[tail++] = e.to;
                }
            }
            for (int v = 0; v < _n; v++) height_count[height[v]]++;
            rebuild_buckets();
            work = 0;
        };

        auto gap = [&](int empty_height) {
            for (int v = 0; v < _n; v++) {
                if (v == s || v == t || height[v] <= empty_height ||
                    height[v] >= _n) {
                    continue;
                }
                height_count[height[v]]--;
                height[v] = unreachable;
                height_count[height[v]]++;
                current[v] = 0;
            }
            rebuild_buckets();
        };

        auto relabel = [&](int v) -> bool {
            int old_height = height[v];
            int new_height = dead;
            work += int(_g[v].size());
            for (const auto& e : _g[v]) {
                if (e.cap != Cap(0)) {
                    new_height = std::min(new_height, height[e.to] + 1);
                }
            }
            height_count[old_height]--;
            height[v] = std::min(new_height, dead);
            height_count[height[v]]++;
            current[v] = 0;
            if (old_height < _n && height_count[old_height] == 0) {
                gap(old_height);
                return true;
            }
            return false;
        };

        auto push = [&](int v, InternalEdge& e) {
            Cap sent = std::min(excess[v], e.cap);
            bool was_zero = excess[e.to] == Cap(0);
            e.cap -= sent;
            _g[e.to][e.rev].cap += sent;
            excess[v] -= sent;
            excess[e.to] += sent;
            if (was_zero) activate(e.to);
        };

        auto discharge = [&](int v) {
            while (excess[v] != Cap(0) && height[v] < dead) {
                if (current[v] == int(_g[v].size())) {
                    if (relabel(v)) return;
                    continue;
                }
                auto& e = _g[v][current[v]];
                work++;
                if (e.cap != Cap(0) && height[v] == height[e.to] + 1) {
                    push(v, e);
                } else {
                    current[v]++;
                }
            }
            activate(v);
        };

        for (auto& e : _g[s]) {
            if (e.to == s || e.cap == Cap(0)) continue;
            Cap sent = e.cap;
            e.cap = Cap(0);
            _g[e.to][e.rev].cap += sent;
            excess[e.to] += sent;
        }
        global_relabel();

        while (highest >= 0) {
            if (bucket_head[highest] == -1) {
                highest--;
                continue;
            }
            int v = bucket_head[highest];
            bucket_head[highest] = next[v];
            if (!active[v] || height[v] != highest) continue;
            active[v] = false;
            discharge(v);
            if (work >= work_limit) global_relabel();
        }
        return excess[t];
    }

   public:
    MaxFlow() : MaxFlow(0) {}

    explicit MaxFlow(int n) : _n(n), _g(n) {
        assert(0 <= n);
    }

    int size() const {
        return _n;
    }

    int edge_count() const {
        return int(_pos.size());
    }

    void reserve_edges(int edge_count) {
        assert(0 <= edge_count);
        _pos.reserve(edge_count);
        if (_n == 0 || edge_count == 0 ||
            2 * std::size_t(edge_count) < std::size_t(_n)) {
            return;
        }
        const std::size_t average_degree =
            (3 * std::size_t(edge_count) + std::size_t(_n) - 1)
            / std::size_t(_n);
        for (auto& edges : _g) edges.reserve(average_degree);
    }

    void reserve_edges(int edge_count, const std::vector<int>& degrees) {
        assert(0 <= edge_count);
        assert(int(degrees.size()) == _n);
        _pos.reserve(edge_count);
        for (int v = 0; v < _n; v++) {
            assert(0 <= degrees[v]);
            _g[v].reserve(degrees[v]);
        }
    }

    int add_edge(int from, int to, Cap cap) {
        assert(0 <= from && from < _n);
        assert(0 <= to && to < _n);
        assert(Cap(0) <= cap);
        int id = int(_pos.size());
        int from_id = int(_g[from].size());
        int to_id = int(_g[to].size());
        if (from == to) to_id++;
        _pos.push_back(Position{from, from_id});
        _g[from].push_back(InternalEdge{to, to_id, cap});
        _g[to].push_back(InternalEdge{from, from_id, Cap(0)});
        return id;
    }

    int add_undirected_edge(int first, int second, Cap cap) {
        static_assert(std::numeric_limits<Cap>::is_signed);
        assert(0 <= first && first < _n);
        assert(0 <= second && second < _n);
        assert(Cap(0) <= cap);
        assert(cap <= std::numeric_limits<Cap>::max() / Cap(2));
        int id = int(_pos.size());
        int first_id = int(_g[first].size());
        int second_id = int(_g[second].size());
        if (first == second) second_id++;
        _pos.push_back(Position{first, ~first_id});
        _g[first].push_back(InternalEdge{second, second_id, cap});
        _g[second].push_back(InternalEdge{first, first_id, cap});
        return id;
    }

    Edge get_edge(int i) const {
        assert(0 <= i && i < int(_pos.size()));
        const auto& position = _pos[i];
        int from = position.from;
        bool undirected = position.edge < 0;
        int idx = undirected ? ~position.edge : position.edge;
        const auto& e = _g[from][idx];
        const auto& re = _g[e.to][e.rev];
        if (undirected) {
            return Edge{
                from,
                e.to,
                (e.cap + re.cap) / Cap(2),
                (re.cap - e.cap) / Cap(2)
            };
        }
        return Edge{from, e.to, e.cap + re.cap, re.cap};
    }

    std::vector<Edge> edges() const {
        std::vector<Edge> result;
        result.reserve(_pos.size());
        for (int i = 0; i < int(_pos.size()); i++) result.push_back(get_edge(i));
        return result;
    }

    void change_edge(int i, Cap new_cap, Cap new_flow) {
        assert(0 <= i && i < int(_pos.size()));
        assert(Cap(0) <= new_cap);
        auto& position = _pos[i];
        int from = position.from;
        bool undirected = position.edge < 0;
        int idx = undirected ? ~position.edge : position.edge;
        auto& e = _g[from][idx];
        auto& re = _g[e.to][e.rev];
        if (undirected) {
            assert(new_cap <= std::numeric_limits<Cap>::max() / Cap(2));
            assert(-new_cap <= new_flow && new_flow <= new_cap);
            e.cap = new_cap - new_flow;
            re.cap = new_cap + new_flow;
        } else {
            assert(Cap(0) <= new_flow && new_flow <= new_cap);
            e.cap = new_cap - new_flow;
            re.cap = new_flow;
        }
    }

    Cap max_flow(int s, int t) {
        assert(0 <= s && s < _n);
        assert(0 <= t && t < _n);
        assert(s != t);
        return highest_label_preflow_push(s, t);
    }

    Cap max_flow_push_relabel(int s, int t) {
        assert(0 <= s && s < _n);
        assert(0 <= t && t < _n);
        assert(s != t);
        return highest_label_preflow_push(s, t);
    }

    Cap max_flow_dinic(int s, int t) {
        return max_flow(s, t, std::numeric_limits<Cap>::max());
    }

    Cap max_flow(int s, int t, Cap flow_limit) {
        assert(0 <= s && s < _n);
        assert(0 <= t && t < _n);
        assert(s != t);

        std::vector<int> work(3 * std::size_t(_n));
        int* level = work.data();
        int* iter = level + _n;
        int* queue = iter + _n;
        auto bfs = [&]() -> bool {
            std::fill(level, level + _n, -1);
            int head = 0;
            int tail = 0;
            level[s] = 0;
            queue[tail++] = s;
            while (head != tail) {
                int v = queue[head++];
                for (const auto& e : _g[v]) {
                    if (level[e.to] != -1 || e.cap == Cap(0)) continue;
                    level[e.to] = level[v] + 1;
                    if (e.to == t) return true;
                    queue[tail++] = e.to;
                }
            }
            return level[t] != -1;
        };

        auto dfs = [&](auto&& self, int v, Cap up) -> Cap {
            if (v == s) return up;
            Cap result = Cap(0);
            const int current_level = level[v];
            auto& edges = _g[v];
            const int edge_count = int(edges.size());
            for (int& i = iter[v]; i < edge_count; i++) {
                auto& e = edges[i];
                if (level[e.to] + 1 != current_level) continue;
                auto& reverse = _g[e.to][e.rev];
                if (reverse.cap == Cap(0)) continue;
                Cap d = self(
                    self,
                    e.to,
                    std::min(up - result, reverse.cap)
                );
                if (d == Cap(0)) continue;
                e.cap += d;
                reverse.cap -= d;
                result += d;
                if (result == up) return result;
            }
            level[v] = _n;
            return result;
        };

        Cap flow = 0;
        while (flow < flow_limit && bfs()) {
            std::fill(iter, iter + _n, 0);
            flow += dfs(dfs, t, flow_limit - flow);
        }
        return flow;
    }

    std::vector<bool> min_cut(int s) const {
        assert(0 <= s && s < _n);
        std::vector<bool> visited(_n, false);
        std::vector<int> queue(_n);
        int head = 0;
        int tail = 0;
        visited[s] = true;
        queue[tail++] = s;
        while (head != tail) {
            int v = queue[head++];
            for (const auto& e : _g[v]) {
                if (e.cap == Cap(0) || visited[e.to]) continue;
                visited[e.to] = true;
                queue[tail++] = e.to;
            }
        }
        return visited;
    }
};

}  // namespace flow
}  // namespace m1une


#line 11 "optimization/project_selection.hpp"

namespace m1une {
namespace opt {

template <class T>
struct ProjectSelectionResult {
    bool feasible;
    T max_gain;
    std::vector<bool> selected;

    bool is_feasible() const {
        return feasible;
    }
};

template <class T>
class ProjectSelection {
    static_assert(std::is_integral_v<T> && std::is_signed_v<T>);

    struct Arc {
        int from;
        int to;
        T cap;
    };

    static constexpr int source = -1;
    static constexpr int sink = -2;

    int _project_count;
    int _vertex_count;
    T _offset = T();
    T _finite_cap_sum = T();
    std::vector<Arc> _arcs;
    std::vector<std::pair<int, int>> _hard_arcs;

    void assert_project(int project) const {
        (void)project;
        assert(0 <= project && project < _project_count);
    }

    void assert_vertex(int vertex) const {
        (void)vertex;
        assert(0 <= vertex && vertex < _vertex_count);
    }

    void add_offset(T value) {
        if (value > T()) {
            assert(_offset <= std::numeric_limits<T>::max() - value);
        } else if (value < T()) {
            assert(_offset >= std::numeric_limits<T>::lowest() - value);
        }
        _offset += value;
    }

    T nonnegative_difference(T large, T small) const {
        assert(small <= large);
        if (small < T()) {
            assert(large <= std::numeric_limits<T>::max() + small);
        }
        return large - small;
    }

    void add_arc(int from, int to, T cap) {
        assert(cap >= T());
        if (from == to) return;
        assert(cap <= std::numeric_limits<T>::max() - _finite_cap_sum);
        _finite_cap_sum += cap;
        _arcs.push_back(Arc{from, to, cap});
    }

    void add_hard_arc(int from, int to) {
        if (from == to) return;
        _hard_arcs.emplace_back(from, to);
    }

    void add_vertex_gain(int vertex, T gain_if_selected, T gain_if_unselected) {
        assert_vertex(vertex);
        if (gain_if_selected >= gain_if_unselected) {
            add_offset(gain_if_selected);
            add_arc(source, vertex,
                    nonnegative_difference(gain_if_selected, gain_if_unselected));
        } else {
            add_offset(gain_if_unselected);
            add_arc(vertex, sink,
                    nonnegative_difference(gain_if_unselected, gain_if_selected));
        }
    }

    int add_auxiliary_vertex() {
        return _vertex_count++;
    }

   public:
    ProjectSelection() : ProjectSelection(0) {}

    explicit ProjectSelection(int project_count)
        : _project_count(project_count), _vertex_count(project_count) {
        assert(project_count >= 0);
    }

    int size() const {
        return _project_count;
    }

    void add_gain(int project, T gain_if_selected) {
        add_gain(project, gain_if_selected, T());
    }

    void add_gain(int project, T gain_if_selected, T gain_if_unselected) {
        assert_project(project);
        add_vertex_gain(project, gain_if_selected, gain_if_unselected);
    }

    void add_penalty(int selected_project, int unselected_project, T penalty) {
        assert_project(selected_project);
        assert_project(unselected_project);
        add_arc(selected_project, unselected_project, penalty);
    }

    void add_penalty_if_different(int project_a, int project_b, T penalty) {
        assert_project(project_a);
        assert_project(project_b);
        add_arc(project_a, project_b, penalty);
        add_arc(project_b, project_a, penalty);
    }

    void add_gain_if_same(int project_a, int project_b, T gain) {
        assert(gain >= T());
        add_offset(gain);
        add_penalty_if_different(project_a, project_b, gain);
    }

    void add_hard_implication(int selected_project, int required_project) {
        assert_project(selected_project);
        assert_project(required_project);
        add_hard_arc(selected_project, required_project);
    }

    void force_selected(int project) {
        assert_project(project);
        add_hard_arc(source, project);
    }

    void force_unselected(int project) {
        assert_project(project);
        add_hard_arc(project, sink);
    }

    void add_gain_if_all_selected(const std::vector<int>& projects, T gain) {
        assert(gain >= T());
        for (int project : projects) assert_project(project);
        if (projects.empty()) {
            add_offset(gain);
            return;
        }
        if (projects.size() == 1) {
            add_vertex_gain(projects[0], gain, T());
            return;
        }
        if (projects.size() == 2) {
            add_vertex_gain(projects[0], gain, T());
            add_arc(projects[0], projects[1], gain);
            return;
        }

        int auxiliary = add_auxiliary_vertex();
        add_vertex_gain(auxiliary, gain, T());
        for (int project : projects) add_hard_arc(auxiliary, project);
    }

    void add_gain_if_all_unselected(const std::vector<int>& projects, T gain) {
        assert(gain >= T());
        for (int project : projects) assert_project(project);
        if (projects.empty()) {
            add_offset(gain);
            return;
        }
        if (projects.size() == 1) {
            add_vertex_gain(projects[0], T(), gain);
            return;
        }
        if (projects.size() == 2) {
            add_vertex_gain(projects[0], T(), gain);
            add_arc(projects[1], projects[0], gain);
            return;
        }

        int auxiliary = add_auxiliary_vertex();
        add_vertex_gain(auxiliary, T(), gain);
        for (int project : projects) add_hard_arc(project, auxiliary);
    }

    ProjectSelectionResult<T> solve() const {
        int s = _vertex_count;
        int t = s + 1;
        flow::MaxFlow<T> max_flow(_vertex_count + 2);

        auto vertex_id = [&](int vertex) {
            if (vertex == source) return s;
            if (vertex == sink) return t;
            return vertex;
        };

        for (const auto& arc : _arcs) {
            max_flow.add_edge(vertex_id(arc.from), vertex_id(arc.to), arc.cap);
        }

        T hard_cap = T();
        if (!_hard_arcs.empty()) {
            assert(_finite_cap_sum < std::numeric_limits<T>::max());
            hard_cap = _finite_cap_sum + T(1);
            for (auto [from, to] : _hard_arcs) {
                max_flow.add_edge(vertex_id(from), vertex_id(to), hard_cap);
            }
        }

        T cut_cost =
            _hard_arcs.empty() ? max_flow.max_flow(s, t) : max_flow.max_flow(s, t, hard_cap);
        ProjectSelectionResult<T> result;
        result.feasible = _hard_arcs.empty() || cut_cost < hard_cap;
        result.max_gain = T();
        result.selected.assign(_project_count, false);
        if (!result.feasible) return result;

        assert(_offset >= std::numeric_limits<T>::lowest() + cut_cost);
        result.max_gain = _offset - cut_cost;
        auto source_side = max_flow.min_cut(s);
        for (int project = 0; project < _project_count; project++) {
            result.selected[project] = source_side[project];
        }
        return result;
    }
};

}  // namespace opt
}  // namespace m1une


#line 12 "optimization/k_project_selection.hpp"

namespace m1une {
namespace opt {

template <class T>
struct KProjectSelectionResult {
    bool feasible;
    T max_gain;
    std::vector<int> values;

    bool is_feasible() const {
        return feasible;
    }
};

template <class T>
class KProjectSelection {
    static_assert(std::is_integral_v<T> && std::is_signed_v<T>);
    static_assert(sizeof(T) <= sizeof(long long));

    using Wide = __int128_t;

    std::vector<int> _value_counts;
    std::vector<int> _first_threshold;
    ProjectSelection<T> _binary;
    T _constant = T();

    static int threshold_count(const std::vector<int>& value_counts) {
        assert(value_counts.size() <=
               std::size_t(std::numeric_limits<int>::max()));
        long long count = 0;
        for (int value_count : value_counts) {
            assert(value_count >= 1);
            count += value_count - 1;
            assert(count <= std::numeric_limits<int>::max());
        }
        return int(count);
    }

    static std::vector<int> repeated_value_counts(
        int project_count,
        int value_count
    ) {
        assert(project_count >= 0);
        assert(value_count >= 1);
        return std::vector<int>(project_count, value_count);
    }

    void assert_project(int project) const {
        (void)project;
        assert(0 <= project && project < size());
    }

    int threshold(int project, int value) const {
        assert_project(project);
        (void)value;
        assert(1 <= value && value < _value_counts[project]);
        return _first_threshold[project] + value - 1;
    }

    static T narrow(Wide value) {
        assert(Wide(std::numeric_limits<T>::lowest()) <= value);
        assert(value <= Wide(std::numeric_limits<T>::max()));
        return T(value);
    }

    void add_constant(T gain) {
        _constant = narrow(Wide(_constant) + gain);
    }

    void add_threshold_gain(int project, int value, Wide gain) {
        if (gain == 0) return;
        _binary.add_gain(threshold(project, value), narrow(gain));
    }

   public:
    KProjectSelection() : KProjectSelection(std::vector<int>()) {}

    explicit KProjectSelection(std::vector<int> value_counts)
        : _value_counts(std::move(value_counts)),
          _first_threshold(_value_counts.size()),
          _binary(threshold_count(_value_counts)) {
        int first = 0;
        for (int project = 0; project < size(); project++) {
            _first_threshold[project] = first;
            first += _value_counts[project] - 1;
        }

        for (int project = 0; project < size(); project++) {
            for (int value = 2; value < _value_counts[project]; value++) {
                _binary.add_hard_implication(
                    threshold(project, value),
                    threshold(project, value - 1)
                );
            }
        }
    }

    KProjectSelection(int project_count, int value_count)
        : KProjectSelection(repeated_value_counts(project_count, value_count)) {}

    int size() const {
        return int(_value_counts.size());
    }

    int value_count(int project) const {
        assert_project(project);
        return _value_counts[project];
    }

    void add_gain(int project, const std::vector<T>& gains) {
        assert_project(project);
        assert(int(gains.size()) == _value_counts[project]);
        add_constant(gains[0]);
        for (int value = 1; value < _value_counts[project]; value++) {
            add_threshold_gain(
                project,
                value,
                Wide(gains[value]) - gains[value - 1]
            );
        }
    }

    void add_gain(
        int project_a,
        int project_b,
        const std::vector<std::vector<T>>& gains
    ) {
        assert_project(project_a);
        assert_project(project_b);
        assert(project_a != project_b);
        const int count_a = _value_counts[project_a];
        const int count_b = _value_counts[project_b];
        assert(int(gains.size()) == count_a);
        for (const auto& row : gains) assert(int(row.size()) == count_b);

        add_constant(gains[0][0]);
        for (int value_a = 1; value_a < count_a; value_a++) {
            add_threshold_gain(
                project_a,
                value_a,
                Wide(gains[value_a][0]) - gains[value_a - 1][0]
            );
        }
        for (int value_b = 1; value_b < count_b; value_b++) {
            add_threshold_gain(
                project_b,
                value_b,
                Wide(gains[0][value_b]) - gains[0][value_b - 1]
            );
        }

        for (int value_a = 1; value_a < count_a; value_a++) {
            for (int value_b = 1; value_b < count_b; value_b++) {
                Wide mixed =
                    Wide(gains[value_a][value_b])
                    - gains[value_a - 1][value_b]
                    - gains[value_a][value_b - 1]
                    + gains[value_a - 1][value_b - 1];
                assert(mixed >= 0);
                T gain = narrow(mixed);
                if (gain == T()) continue;
                int threshold_a = threshold(project_a, value_a);
                int threshold_b = threshold(project_b, value_b);
                _binary.add_gain(threshold_a, gain);
                _binary.add_penalty(threshold_a, threshold_b, gain);
            }
        }
    }

    void force_value(int project, int value) {
        assert_project(project);
        assert(0 <= value && value < _value_counts[project]);
        force_value_at_least(project, value);
        force_value_at_most(project, value);
    }

    void force_value_at_least(int project, int lower_bound) {
        assert_project(project);
        assert(0 <= lower_bound && lower_bound < _value_counts[project]);
        if (lower_bound > 0) {
            _binary.force_selected(threshold(project, lower_bound));
        }
    }

    void force_value_at_most(int project, int upper_bound) {
        assert_project(project);
        assert(0 <= upper_bound && upper_bound < _value_counts[project]);
        if (upper_bound + 1 < _value_counts[project]) {
            _binary.force_unselected(threshold(project, upper_bound + 1));
        }
    }

    KProjectSelectionResult<T> solve() const {
        auto binary_result = _binary.solve();
        KProjectSelectionResult<T> result;
        result.feasible = binary_result.feasible;
        result.max_gain = T();
        result.values.assign(size(), 0);
        if (!result.feasible) return result;

        result.max_gain = narrow(Wide(_constant) + binary_result.max_gain);
        for (int project = 0; project < size(); project++) {
            for (int value = 1; value < _value_counts[project]; value++) {
                if (!binary_result.selected[threshold(project, value)]) break;
                result.values[project] = value;
            }
        }
        return result;
    }
};

}  // namespace opt
}  // namespace m1une
Back to top page