提出 #61474030

---

ソースコード 拡げる

#include <bits/stdc++.h>

const int INF = 1u << 30u; // 1,073,741,824

// capacity scaling + dinic
// O(EV log U)
template <typename T>
class Dinic {
public:
    struct Edge {
        const int from;
        const int to;
        T flow;
        const T cap;
        const int rev;

        Edge(const int from, const int to, const T flow, const T cap, const int rev) :
            from(from), to(to), flow(flow), cap(cap), rev(rev) {
            assert(this->cap >= 0);
        }

        T residual_capacity() const { return this->cap - this->flow; }
    };

    int num_nodes;
    int num_edges;
    std::vector<std::vector<Edge>> graph;
    std::vector<int> level;
    std::vector<int> current_edge;
    std::vector<std::pair<int, int>> edge_id_memo;

    Dinic() : num_nodes(0), num_edges(0) {}

    int add_node() {
        this->add_nodes(1);
        return this->num_nodes - 1;
    }

    std::vector<int> add_nodes(const int num) {
        std::vector<int> nodes(num);
        std::iota(nodes.begin(), nodes.end(), this->num_nodes);
        this->num_nodes += num;
        this->graph.resize(this->num_nodes);
        return nodes;
    }

    int add_directed_edge(const int from, const int to, const T cap) {
        assert(0 <= from and from < this->num_nodes and 0 <= to and to < this->num_nodes);
        assert(cap >= 0);
        this->graph[from].emplace_back(from, to, 0, cap, static_cast<int>(graph[to].size()));
        this->graph[to].emplace_back(to, from, cap, cap, static_cast<int>(graph[from].size()) - 1);
        this->edge_id_memo.emplace_back(from, static_cast<int>(this->graph[from].size()) - 1);
        return this->num_edges++;
    }

    Edge get_edge(const int edge_id) {
        const auto [u, i] = this->edge_id_memo[edge_id];
        return this->graph[u][i];
    }

    T solve(const int source, const int sink) {
        assert(source < this->num_nodes and sink < this->num_nodes);
        this->level.resize(this->num_nodes);
        this->current_edge.resize(this->num_nodes);

        T max_capacity = 0;
        for (int u = 0; u < this->num_nodes; ++u) {
            for (const auto& e : this->graph[u]) {
                max_capacity = std::max(max_capacity, e.cap);
            }
        }
        T delta = 1;
        while (delta <= max_capacity) {
            delta *= 2;
        }
        delta /= 2;

        T upper = 0;
        for (const auto& e : this->graph[source]) {
            upper += e.cap;
        }

        T flow = 0;
        while (delta > 0) {
            // solve maximum flow in delta-residual network
            while (true) {
                this->bfs(source, sink, delta);

                // no s-t path
                if (this->level[source] >= this->num_nodes) {
                    break;
                }

                fill(this->current_edge.begin(), this->current_edge.end(), 0);
                flow += dfs(source, sink, upper, delta);
            }
            delta /= 2;
        }

        return flow;
    }

    std::vector<bool> minimum_cut(const int source) {
        std::vector<bool> visited(this->num_nodes);
        std::queue<int> que;
        que.emplace(source);
        visited[source] = true;

        while (not que.empty()) {
            const auto u = que.front();
            que.pop();

            for (const auto& e : this->graph[u]) {
                if (not visited[e.to] and e.residual_capacity() != 0) {
                    visited[e.to] = true;
                    que.emplace(e.to);
                }
            }
        }

        return visited;
    }

private:
    void bfs(int source, int sink, T delta) {
        fill(this->level.begin(), this->level.end(), this->num_nodes);
        std::queue<int> que;
        this->level[sink] = 0;
        que.push(sink);
        while (not que.empty()) {
            auto v = que.front();
            que.pop();

            for (const auto& e : this->graph[v]) {
                // check e.to -> v
                if (e.flow >= delta and level[e.to] == this->num_nodes) {
                    this->level[e.to] = this->level[v] + 1;
                    if (e.to != source) {
                        que.push(e.to);
                    }
                }
            }
        }
    }

    T dfs(const int u, const int sink, T upper, T delta) {
        if (u == sink) {
            return upper;
        }

        T flow = 0;
        for (int& i = this->current_edge[u]; i < static_cast<int>(this->graph[u].size()); ++i) {
            auto& e = this->graph[u][i];
            const auto residual_capacity = e.residual_capacity();
            if (residual_capacity >= delta and this->level[u] > this->level[e.to]) {
                const auto d = dfs(e.to, sink, std::min(upper - flow, residual_capacity), delta);
                // update flow
                e.flow += d;
                this->graph[e.to][e.rev].flow -= d;

                flow += d;
                if (flow == upper or d == 0) {
                    return flow;
                }
            }
        }
        this->level[u] = this->num_nodes;

        return flow;
    }
};

// Quadratic pseudo-Boolean optimization
// Reference: Minimizing Nonsubmodular Functions: A Review, DOI: 10.1109/TPAMI.2007.1031
// 関数が劣モジュラのとき最適解を求めることができる
template <class COST>
class QPBO {
    int num_variables;
    std::vector<std::array<COST, 2>> unary_costs;
    std::map<std::pair<int, int>, std::array<COST, 4>> pair_wise_costs;
    Dinic<COST> dinic;
    std::vector<int> labels;
    std::vector<int> xs, ys;
    int source, sink;

public:
    QPBO() : num_variables(0), source(-1), sink(-1) {}

    int add_variable() {
        this->add_variables(1);
        return this->num_variables - 1;
    }

    std::vector<int> add_variables(const int num) {
        std::vector<int> nodes(num);
        std::iota(nodes.begin(), nodes.end(), this->num_variables);
        this->num_variables += num;
        this->unary_costs.resize(this->num_variables);
        this->labels.resize(this->num_variables, -1);
        return nodes;
    }

    // f(i = b) = cost
    void add_unary_cost(const int i, const int b, const COST cost) {
        assert(0 <= i and i < this->num_variables);
        assert(0 == b or b == 1);
        this->unary_costs[i][b] += cost;
    }

    // f(i = 0) = cost_0, f(i = 1) = cost_1
    void add_unary_cost_all(const int i, const COST cost_0, const COST cost_1) {
        assert(0 <= i and i < this->num_variables);
        this->unary_costs[i][0b0] += cost_0;
        this->unary_costs[i][0b1] += cost_1;
    }

    // f(i = b1, j = b2) = cost
    void add_pairwise_cost(const int i, const int j, const int b1, const int b2, const COST cost) {
        assert(0 <= i and i < this->num_variables and 0 <= j and j < this->num_variables);
        assert((0 == b1 or b1 == 1) and (0 == b2 or b2 == 1));
        this->pair_wise_costs[{i, j}][static_cast<int>(b1) << 1 | b2] += cost;
    }

    // f(i = 0, j = 0) = cost_00, f(i = 0, j = 1) = cost_01, f(i = 1, j = 0) = cost_10, f(i = 1, j = 1) = cost_11
    void add_pairwise_cost_all(const int i, const int j, const COST cost_00, const COST cost_01, const COST cost_10,
                               const COST cost_11) {
        assert(0 <= i and i < this->num_variables and 0 <= j and j < this->num_variables);
        this->pair_wise_costs[{i, j}][0b00] += cost_00;
        this->pair_wise_costs[{i, j}][0b01] += cost_01;
        this->pair_wise_costs[{i, j}][0b10] += cost_10;
        this->pair_wise_costs[{i, j}][0b11] += cost_11;
    }

    COST solve() {
        const auto offset = this->re_parameterization();

        this->xs = this->dinic.add_nodes(this->num_variables);
        this->ys = this->dinic.add_nodes(this->num_variables);
        this->source = this->dinic.add_node();
        this->sink = this->dinic.add_node();

        std::vector<int> tmp_edges;
        for (int p = 0; p < this->num_variables; ++p) {
            const auto& cost = this->unary_costs[p];
            assert(std::min(cost[0b0], cost[0b1]) == 0);
            if (cost[0b0] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->xs[p], sink, cost[0b0]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(source, this->ys[p], cost[0b0]));
            }
            if (cost[0b1] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(source, this->xs[p], cost[0b1]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->ys[p], sink, cost[0b1]));
            }
        }

        for (const auto& [key, cost] : this->pair_wise_costs) {
            const auto [p, q] = key;
            assert(std::min(cost[0b00], cost[0b10]) == 0);
            assert(std::min(cost[0b01], cost[0b11]) == 0);

            if (cost[0b00] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->xs[p], this->ys[q], cost[0b00]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->xs[q], this->ys[p], cost[0b00]));
            }
            if (cost[0b01] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->xs[p], this->xs[q], cost[0b01]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->ys[q], this->ys[p], cost[0b01]));
            }
            if (cost[0b10] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->xs[q], this->xs[p], cost[0b10]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->ys[p], this->ys[q], cost[0b10]));
            }
            if (cost[0b11] != 0) {
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->ys[q], this->xs[p], cost[0b11]));
                tmp_edges.emplace_back(this->dinic.add_directed_edge(this->ys[p], this->xs[q], cost[0b11]));
            }
        }

        return this->dinic.solve(this->source, this->sink) / 2 + offset;
    }

private:
    COST re_parameterization() {
        for (auto& [key, cost] : this->pair_wise_costs) {
            const auto [p, q] = key;
            for (int b = 0; b <= 1; ++b) {
                const auto delta = std::min(cost[0b00 | b], cost[0b10 | b]);
                cost[0b00 | b] -= delta;
                cost[0b10 | b] -= delta;
                this->unary_costs[q][b] += delta;
            }
        }

        COST offset = 0;
        for (int p = 0; p < this->num_variables; ++p) {
            auto& cost = this->unary_costs[p];
            const auto delta = std::min(cost[0b0], cost[0b1]);
            cost[0b0] -= delta;
            cost[0b1] -= delta;
            offset += delta;
        }

        return offset;
    }
};


using namespace std;

int main() {
    cin.tie(nullptr);
    ios::sync_with_stdio(false);

    int N;
    cin >> N;
    vector<int> A(N);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
    }

    QPBO<long long> solver;
    const auto xs = solver.add_variables(N);

    for (int i = 0; i < N; ++i) {
        solver.add_unary_cost(xs[i], 0, -A[i]);
    }

    for (int i = 0; i < N; ++i) {
        for (int j = i + 1; j < N; ++j) {
            if ((j + 1) % (i + 1) == 0) {
                solver.add_pairwise_cost(xs[i], xs[j], 1, 0, INF);
            }
        }
    }
    cout << -solver.solve() << endl;

    return 0;
}

提出情報

ジャッジ結果