#690755 (C++17) No.1324 Approximate the Matrix

• 問題: No.1324 Approximate the Matrix
• ユーザー: hitonanode
• 提出日時: 2021-08-19 00:22:26
• 言語: C++17 (gcc 15.2.0 + boost 1.89.0)
• 結果: AC
• 実行時間: 66 ms / 2,000 ms
• コード長: 15,127 bytes
• コンパイル時間: 1,687 ms
• コンパイル使用メモリ: 145,140 KB
• 最終ジャッジ日時: 2025-01-23 22:58:54
• ジャッジサーバーID （参考情報）: judge5 / judge1

ファイルパターン	結果
sample	AC * 3
other	AC * 42

ソースコード

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <chrono>
#include <cmath>
#include <deque>
#include <forward_list>
#include <fstream>
#include <functional>
#include <iomanip>
#include <ios>
#include <iostream>
#include <limits>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
using namespace std;
using lint = long long;
using pint = pair<int, int>;
using plint = pair<lint, lint>;
struct fast_ios {
    fast_ios() { cin.tie(nullptr), ios::sync_with_stdio(false), cout << fixed << setprecision(20); };
} fast_ios_;
#define ALL(x) (x).begin(), (x).end()
#define FOR(i, begin, end) for (int i = (begin), i##_end_ = (end); i < i##_end_; i++)
#define IFOR(i, begin, end) for (int i = (end)-1, i##_begin_ = (begin); i >= i##_begin_; i--)
#define REP(i, n) FOR(i, 0, n)
#define IREP(i, n) IFOR(i, 0, n)
template <typename T, typename V> void ndarray(vector<T> &vec, const V &val, int len) {
    vec.assign(len, val);
}
template <typename T, typename V, typename... Args>
void ndarray(vector<T> &vec, const V &val, int len, Args... args) {
    vec.resize(len), for_each(begin(vec), end(vec), [&](T &v) { ndarray(v, val, args...); });
}
template <typename T> bool chmax(T &m, const T q) {
    if (m < q) {
        m = q;
        return true;
    } else
        return false;
}
template <typename T> bool chmin(T &m, const T q) {
    if (m > q) {
        m = q;
        return true;
    } else
        return false;
}
int floor_lg(long long x) { return x <= 0 ? -1 : 63 - __builtin_clzll(x); }
template <typename T1, typename T2> pair<T1, T2> operator+(const pair<T1, T2> &l, const pair<T1, T2> &r) {
    return make_pair(l.first + r.first, l.second + r.second);
}
template <typename T1, typename T2> pair<T1, T2> operator-(const pair<T1, T2> &l, const pair<T1, T2> &r) {
    return make_pair(l.first - r.first, l.second - r.second);
}
template <typename T> vector<T> sort_unique(vector<T> vec) {
    sort(vec.begin(), vec.end()), vec.erase(unique(vec.begin(), vec.end()), vec.end());
    return vec;
}
template <typename T> int arglb(const std::vector<T> &v, const T &x) {
    return std::distance(v.begin(), std::lower_bound(v.begin(), v.end(), x));
}
template <typename T> int argub(const std::vector<T> &v, const T &x) {
    return std::distance(v.begin(), std::lower_bound(v.begin(), v.end(), x));
}
template <typename T> istream &operator>>(istream &is, vector<T> &vec) {
    for (auto &v : vec) is >> v;
    return is;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
    os << '[';
    for (auto v : vec) os << v << ',';
    os << ']';
    return os;
}
template <typename T, size_t sz> ostream &operator<<(ostream &os, const array<T, sz> &arr) {
    os << '[';
    for (auto v : arr) os << v << ',';
    os << ']';
    return os;
}
#if __cplusplus >= 201703L
template <typename... T> istream &operator>>(istream &is, tuple<T...> &tpl) {
    std::apply([&is](auto &&...args) { ((is >> args), ...); }, tpl);
    return is;
}
template <typename... T> ostream &operator<<(ostream &os, const tuple<T...> &tpl) {
    std::apply([&os](auto &&...args) { ((os << args << ','), ...); }, tpl);
    return os;
}
#endif
template <typename T> ostream &operator<<(ostream &os, const deque<T> &vec) {
    os << "deq[";
    for (auto v : vec) os << v << ',';
    os << ']';
    return os;
}
template <typename T> ostream &operator<<(ostream &os, const set<T> &vec) {
    os << '{';
    for (auto v : vec) os << v << ',';
    os << '}';
    return os;
}
template <typename T, typename TH> ostream &operator<<(ostream &os, const unordered_set<T, TH> &vec) {
    os << '{';
    for (auto v : vec) os << v << ',';
    os << '}';
    return os;
}
template <typename T> ostream &operator<<(ostream &os, const multiset<T> &vec) {
    os << '{';
    for (auto v : vec) os << v << ',';
    os << '}';
    return os;
}
template <typename T> ostream &operator<<(ostream &os, const unordered_multiset<T> &vec) {
    os << '{';
    for (auto v : vec) os << v << ',';
    os << '}';
    return os;
}
template <typename T1, typename T2> ostream &operator<<(ostream &os, const pair<T1, T2> &pa) {
    os << '(' << pa.first << ',' << pa.second << ')';
    return os;
}
template <typename TK, typename TV> ostream &operator<<(ostream &os, const map<TK, TV> &mp) {
    os << '{';
    for (auto v : mp) os << v.first << "=>" << v.second << ',';
    os << '}';
    return os;
}
template <typename TK, typename TV, typename TH>
ostream &operator<<(ostream &os, const unordered_map<TK, TV, TH> &mp) {
    os << '{';
    for (auto v : mp) os << v.first << "=>" << v.second << ',';
    os << '}';
    return os;
}
#ifdef HITONANODE_LOCAL
const string COLOR_RESET = "\033[0m", BRIGHT_GREEN = "\033[1;32m", BRIGHT_RED = "\033[1;31m", BRIGHT_CYAN = "\033[1;36m", NORMAL_CROSSED = "\033[0;9;37m", RED_BACKGROUND = "\033[1;41m", NORMAL_FAINT = "\033[0;2m";
#define dbg(x)                                                                                               \
    cerr << BRIGHT_CYAN << #x << COLOR_RESET << " = " << (x) << NORMAL_FAINT << " (L" << __LINE__ << ") " << __FILE__ << COLOR_RESET << endl
#else
#define dbg(x) (x)
#endif

/*
MinCostFlow: Minimum-cost flow problem solver WITH NO NEGATIVE CYCLE (just negative cost edge is allowed)
Verified by SRM 770 Div1 Medium <https://community.topcoder.com/stat?c=problem_statement&pm=15702>
*/
template <typename CAP = long long, typename COST = long long> struct MinCostFlow {
    const COST INF_COST = std::numeric_limits<COST>::max() / 2;
    struct edge {
        int to, rev;
        CAP cap;
        COST cost;
        friend std::ostream &operator<<(std::ostream &os, const edge &e) {
            os << '(' << e.to << ',' << e.rev << ',' << e.cap << ',' << e.cost << ')';
            return os;
        }
    };
    int V;
    std::vector<std::vector<edge>> g;
    std::vector<COST> dist;
    std::vector<int> prevv, preve;
    std::vector<COST> dual; // dual[V]: potential
    std::vector<std::pair<int, int>> pos;

    bool _calc_distance_bellman_ford(int s) { // O(VE), able to detect negative cycle
        dist.assign(V, INF_COST);
        dist[s] = 0;
        bool upd = true;
        int cnt = 0;
        while (upd) {
            upd = false;
            cnt++;
            if (cnt > V) return false; // Negative cycle existence
            for (int v = 0; v < V; v++)
                if (dist[v] != INF_COST) {
                    for (int i = 0; i < (int)g[v].size(); i++) {
                        edge &e = g[v][i];
                        COST c = dist[v] + e.cost + dual[v] - dual[e.to];
                        if (e.cap > 0 and dist[e.to] > c) {
                            dist[e.to] = c, prevv[e.to] = v, preve[e.to] = i;
                            upd = true;
                        }
                    }
                }
        }
        return true;
    }

    bool _calc_distance_dijkstra(int s) { // O(ElogV)
        dist.assign(V, INF_COST);
        dist[s] = 0;
        using P = std::pair<COST, int>;
        std::priority_queue<P, std::vector<P>, std::greater<P>> q;
        q.emplace(0, s);
        while (!q.empty()) {
            P p = q.top();
            q.pop();
            int v = p.second;
            if (dist[v] < p.first) continue;
            for (int i = 0; i < (int)g[v].size(); i++) {
                edge &e = g[v][i];
                COST c = dist[v] + e.cost + dual[v] - dual[e.to];
                if (e.cap > 0 and dist[e.to] > c) {
                    dist[e.to] = c, prevv[e.to] = v, preve[e.to] = i;
                    q.emplace(dist[e.to], e.to);
                }
            }
        }
        return true;
    }

    MinCostFlow(int V = 0) : V(V), g(V) {}

    void add_edge(int from, int to, CAP cap, COST cost) {
        assert(0 <= from and from < V);
        assert(0 <= to and to < V);
        pos.emplace_back(from, g[from].size());
        g[from].emplace_back(edge{to, (int)g[to].size() + (from == to), cap, cost});
        g[to].emplace_back(edge{from, (int)g[from].size() - 1, (CAP)0, -cost});
    }

    std::pair<CAP, COST> flow(int s, int t, const CAP &f) {
        /*
        Flush amount of `f` from `s` to `t` using the Dijkstra's algorithm
        works for graph with no negative cycles (negative cost edges are allowed)
        retval: (flow, cost)
        */
        COST ret = 0;
        dual.assign(V, 0);
        prevv.assign(V, -1);
        preve.assign(V, -1);
        CAP frem = f;
        while (frem > 0) {
            _calc_distance_dijkstra(s);
            if (dist[t] == INF_COST) break;
            for (int v = 0; v < V; v++) dual[v] = std::min(dual[v] + dist[v], INF_COST);
            CAP d = frem;
            for (int v = t; v != s; v = prevv[v]) { d = std::min(d, g[prevv[v]][preve[v]].cap); }
            frem -= d;
            ret += d * dual[t];
            for (int v = t; v != s; v = prevv[v]) {
                edge &e = g[prevv[v]][preve[v]];
                e.cap -= d;
                g[v][e.rev].cap += d;
            }
        }
        return std::make_pair(f - frem, ret);
    }

    friend std::ostream &operator<<(std::ostream &os, const MinCostFlow &mcf) {
        os << "[MinCostFlow]V=" << mcf.V << ":";
        for (int i = 0; i < (int)mcf.g.size(); i++)
            for (auto &e : mcf.g[i]) {
                os << "\n" << i << "->" << e.to << ": cap=" << e.cap << ", cost=" << e.cost;
            }
        return os;
    }
};

// https://people.orie.cornell.edu/dpw/orie633/
template <class Cap, class Cost, int SCALING = 1, int REFINEMENT_ITER = 20> struct mcf_costscaling {
    mcf_costscaling() = default;
    mcf_costscaling(int n) : _n(n), to(n), b(n), p(n) {}

    int _n;
    std::vector<Cap> cap;
    std::vector<Cost> cost;
    std::vector<int> opposite;
    std::vector<std::vector<int>> to;
    std::vector<Cap> b;
    std::vector<Cost> p;

    int add_edge(int from_, int to_, Cap cap_, Cost cost_) {
        assert(0 <= from_ and from_ < _n);
        assert(0 <= to_ and to_ < _n);
        assert(0 <= cap_);
        cost_ *= (_n + 1);
        const int e = int(cap.size());
        to[from_].push_back(e);
        cap.push_back(cap_);
        cost.push_back(cost_);
        opposite.push_back(to_);

        to[to_].push_back(e + 1);
        cap.push_back(0);
        cost.push_back(-cost_);
        opposite.push_back(from_);
        return e / 2;
    }
    void add_supply(int v, Cap supply) { b[v] += supply; }
    void add_demand(int v, Cap demand) { add_supply(v, -demand); }

    template <typename RetCost = Cost> RetCost solve() {
        Cost eps = 1;
        std::vector<int> que;
        for (const auto c : cost) {
            while (eps <= -c) eps <<= SCALING;
        }
        for (; eps >>= SCALING;) {
            auto no_admissible_cycle = [&]() -> bool {
                for (int i = 0; i < _n; i++) {
                    if (b[i]) return false;
                }
                std::vector<Cost> pp = p;
                for (int iter = 0; iter < REFINEMENT_ITER; iter++) {
                    bool flg = false;
                    for (int e = 0; e < int(cap.size()); e++) {
                        if (!cap[e]) continue;
                        int i = opposite[e ^ 1], j = opposite[e];
                        if (pp[j] > pp[i] + cost[e] + eps) pp[j] = pp[i] + cost[e] + eps, flg = true;
                    }
                    if (!flg) return p = pp, true;
                }
                return false;
            };
            if (no_admissible_cycle()) continue; // Refine

            for (int e = 0; e < int(cap.size()); e++) {
                const int i = opposite[e ^ 1], j = opposite[e];
                const Cost cp_ij = cost[e] + p[i] - p[j];
                if (cap[e] and cp_ij < 0) b[i] -= cap[e], b[j] += cap[e], cap[e ^ 1] += cap[e], cap[e] = 0;
            }
            que.clear();
            int qh = 0;
            for (int i = 0; i < _n; i++) {
                if (b[i] > 0) que.push_back(i);
            }
            std::vector<int> iters(_n);
            while (qh < int(que.size())) {
                const int i = que[qh++];
                for (; iters[i] < int(to[i].size()) and b[i]; ++iters[i]) { // Push
                    int e = to[i][iters[i]];
                    if (!cap[e]) continue;
                    int j = opposite[e];
                    Cost cp_ij = cost[e] + p[i] - p[j];
                    if (cp_ij >= 0) continue;
                    Cap f = b[i] > cap[e] ? cap[e] : b[i];
                    if (b[j] <= 0 and b[j] + f > 0) que.push_back(j);
                    b[i] -= f, b[j] += f, cap[e] -= f, cap[e ^ 1] += f;
                }

                if (b[i] > 0) { // Relabel
                    bool flg = false;
                    for (int e : to[i]) {
                        if (!cap[e]) continue;
                        Cost x = p[opposite[e]] - cost[e] - eps;
                        if (!flg or x > p[i]) flg = true, p[i] = x;
                    }
                    que.push_back(i), iters[i] = 0;
                }
            }
        }
        RetCost ret = 0;
        for (int e = 0; e < int(cap.size()); e += 2) ret += RetCost(cost[e]) * cap[e ^ 1];
        return ret / (_n + 1);
    }
    std::vector<Cost> potential() const {
        std::vector<Cost> ret = p, c0 = cost;
        for (auto &x : ret) x /= (_n + 1);
        for (auto &x : c0) x /= (_n + 1);
        while (true) {
            bool flg = false;
            for (int i = 0; i < _n; i++) {
                for (const auto e : to[i]) {
                    if (!cap[e]) continue;
                    int j = opposite[e];
                    auto y = ret[i] + c0[e];
                    if (ret[j] > y) ret[j] = y, flg = true;
                }
            }
            if (!flg) break;
        }
        return ret;
    }
    struct edge {
        int from, to;
        Cap cap, flow;
        Cost cost;
    };
    edge get_edge(int e) const {
        int m = cap.size() / 2;
        assert(e >= 0 and e < m);
        return {opposite[e * 2 + 1], opposite[e * 2], cap[e * 2] + cap[e * 2 + 1], cap[e * 2 + 1], cost[e * 2] / (_n + 1)};
    }
    std::vector<edge> edges() const {
        int m = cap.size() / 2;
        std::vector<edge> result(m);
        for (int i = 0; i < m; i++) result[i] = get_edge(i);
        return result;
    }
};

int main() {
    int N, K;
    cin >> N >> K;
    vector<int> A(N), B(N);
    vector P(N, vector<int>(N));
    cin >> A >> B >> P;

    const int gs = N * 2, gt = gs + 1;
    mcf_costscaling<int, lint> graph(gt + 1);

    lint ret = 0;
    REP(i, N) REP(j, N) {
        ret += P[i][j] * P[i][j];
        REP(a, A[i]) graph.add_edge(i, j + N, 1, 2 * (a - P[i][j]) + 1);
    }

    REP(i, N) {
        graph.add_edge(gs, i, A[i], 0);
        graph.add_edge(i + N, gt, B[i], 0);
    }
    graph.add_supply(gs, K);
    graph.add_demand(gt, K);
    cout << ret + graph.solve() << '\n';
}