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

• 問題: No.1324 Approximate the Matrix
• ユーザー: theory_and_me
• 提出日時: 2020-11-30 21:33:39
• 言語: C++17 (gcc 12.3.0 + boost 1.83.0)
• 結果: CE (最新)
  AC
  (最初)
• 実行時間: -
• コード長: 9,443 bytes
• コンパイル時間: 710 ms
• コンパイル使用メモリ: 81,768 KB
• 最終ジャッジ日時: 2024-11-14 23:56:12
• 合計ジャッジ時間: 1,087 ms
• ジャッジサーバーID （参考情報）: judge1 / judge5

コンパイルメッセージ

main.cpp:15:14: error: 'function' in namespace 'std' does not name a template type
   15 |         std::function<Cost(Flow)> F; // must be discrete convex
      |              ^~~~~~~~
main.cpp:7:1: note: 'std::function' is defined in header '<functional>'; did you forget to '#include <functional>'?
    6 | #include <iostream>
  +++ |+#include <functional>
    7 | 
main.cpp:36:36: error: 'std::function' has not been declared
   36 |     int add_edge(int from, int to, std::function<Cost(Flow)> F, Flow cap, Flow flow){
      |                                    ^~~
main.cpp:36:49: error: expected ',' or '...' before '<' token
   36 |     int add_edge(int from, int to, std::function<Cost(Flow)> F, Flow cap, Flow flow){
      |                                                 ^
main.cpp: In member function 'int mccf_graph<Flow, Cost>::add_edge(int, int, int)':
main.cpp:41:50: error: 'F' was not declared in this scope
   41 |         E.push_back(edge_with_function{from, to, F, cap, flow});
      |                                                  ^
main.cpp:41:53: error: 'cap' was not declared in this scope
   41 |         E.push_back(edge_with_function{from, to, F, cap, flow});
      |                                                     ^~~
main.cpp:41:58: error: 'flow' was not declared in this scope; did you mean 'Flow'?
   41 |         E.push_back(edge_with_function{from, to, F, cap, flow});
      |                                                          ^~~~
      |                                                          Flow
main.cpp: In function 'int main()':
main.cpp:308:23: error: no matching function for call to 'mccf_graph<long int, long int>::add_edge(int&, int, main()::<lambda(int64_t)>&, __gnu_cxx::__alloc_traits<std::allocator<int>, int>::value_type&, int)'
  308 |             G.add_edge(i, N+j, F, A[i], 0);
      |             ~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~
main.cpp:36:9: note: candidate: 'int mccf_graph<Flow, Cost>::add_edge(int, int, int) [with Flow = lo

ソースコード

#include <algorithm>
#include <cassert>
#include <limits>
#include <queue>
#include <vector>
#include <iostream>

template<class Flow, class Cost> struct mccf_graph{

    int n;

    // input edges
    struct edge_with_function{
        int from, to;
        std::function<Cost(Flow)> F; // must be discrete convex
        Flow cap, flow;
    };

    // edges on the residual graph
    struct edge{
        int to, rev, id;
        bool available;
        Cost cost;
        bool is_rev;
    };

    std::vector<std::vector<edge>> g;
    std::vector<edge_with_function> E;
    std::vector<Cost> dual;

    mccf_graph() {}
    mccf_graph(int n) : n(n), g(n){
        dual.resize(n, 0);
    }

    int add_edge(int from, int to, std::function<Cost(Flow)> F, Flow cap, Flow flow){
        assert(0 <= from and from < n);
        assert(0 <= to and to < n);
        // assert(cap > 0);
        int m = int(E.size());
        E.push_back(edge_with_function{from, to, F, cap, flow});
        return m;
    }

    Cost calc_objective(){
        Cost obj = 0;
        for(auto &e: E){
            obj += e.F(e.flow);
        }
        return obj;
    }
    
    /* for debug
    void print_g(){
        for(int i=0;i<n;i++){
            for(edge &e: g[i]){
                std::cerr << "to " << e.to;
                std::cerr << " rev " << e.rev;
                std::cerr << " id " << e.id;
                std::cerr << " avail " << (int)e.available;
                std::cerr << " cost " << e.cost;
                std::cerr << " r_cost " << e.cost - dual[e.to] + dual[i];
                std::cerr << " is_rev " << (int)e.is_rev;
                std::cerr << "\n";
            }
        }
        std::cerr << "\n";
    }

    void print_E(){
        for(auto &e: E){
            std::cerr << "from " << e.from;
            std::cerr << " to " << e.to;
            std::cerr << " cap " << e.cap;
            std::cerr << " flow " << e.flow;
            std::cerr << "\n";
        }
        std::cerr << "\n";
    }
    for debug */ 

    void solve_CS(std::vector<Flow> b){
        // b is copied, not referenced

        assert((int)b.size() == n);
        std::vector<Cost> dist(n);
        std::vector<int> pv(n), pe(n);
        std::vector<bool> vis(n);

        Flow cap_max = 0;
        for(const auto &e: E) cap_max = std::max(cap_max, e.cap);

        Flow delta = 1;
        while(1){
            if(delta * 2 > cap_max) break;
            delta *= 2;
        }
        Flow init_delta = delta;

        // construct a delta-residual graph
        int m = E.size();
        for(int i=0;i<m;i++){
            auto &e = E[i];
            // if(!e.cap) continue;

            if(delta > e.cap){
                g[e.from].push_back(edge{e.to, (int)g[e.to].size(), i, false, 0, false});
            }else{
                g[e.from].push_back(edge{e.to, (int)g[e.to].size(), i, true, e.F(delta) - e.F(0), false});
            }

            g[e.to].push_back(edge{e.from, (int)g[e.from].size()-1, i, false, 0, true});
        }

        auto check_RCO = [&](int e_id, Flow x, Flow delta)->bool{
            // check edge-wise reduced cost optimality
            auto &e = E[e_id];
            if(x < 0 or x > e.cap) return false;

            bool ok = true;
            if(x + delta <= e.cap){
                Cost reduced_cost = (init_delta/delta)*(e.F(x+delta) - e.F(x)) - dual[e.to] + dual[e.from];
                if(reduced_cost < 0) ok = false;
            }
            if(x - delta >= 0){
                Cost reduced_cost = (init_delta/delta)*(e.F(x-delta) - e.F(x)) - dual[e.from] + dual[e.to];
                if(reduced_cost < 0) ok = false;
            }

            return ok;
        };

        auto delta_update_edge = [&](edge &e, Flow delta){
            auto &e_func = E[e.id];
            if(e.is_rev){
                if(e_func.flow - delta < 0){
                    e.available = false;
                    e.cost = 0;
                }else{
                    e.available = true;
                    e.cost = e_func.F(e_func.flow-delta) - e_func.F(e_func.flow);
                    e.cost *= (init_delta/delta);
                }
            }else{
                if(e_func.flow + delta > e_func.cap){
                    e.available = false;
                    e.cost = 0;
                }else{
                    e.available = true;
                    e.cost = e_func.F(e_func.flow+delta) - e_func.F(e_func.flow);
                    e.cost *= (init_delta/delta);
                }
            }
            return;
        };

        auto saturate = [&](Flow delta){
            // modify edge_with_functions
            for(int i=0;i<m;i++){
                if(check_RCO(i, E[i].flow+delta, delta)){
                    E[i].flow += delta;
                    b[E[i].from] -= delta;
                    b[E[i].to] += delta;
                }
                if(check_RCO(i, E[i].flow-delta, delta)){
                    E[i].flow -= delta;
                    b[E[i].from] += delta;
                    b[E[i].to] -= delta;
                }
                assert(check_RCO(i, E[i].flow, delta));
            }

            // modify edges
            for(int i=0;i<n;i++){
                for(auto &e: g[i]){
                    delta_update_edge(e, delta);
                }
            }
            return;
        };

        auto delta_scaling_dual = [&](Flow delta){
            // return deltat-sink node t
            // if there is no delta-source or delta-sink node, return -1

            std::fill(dist.begin(), dist.end(),
                      std::numeric_limits<Cost>::max());
            std::fill(pv.begin(), pv.end(), -1);
            std::fill(pe.begin(), pe.end(), -1);
            std::fill(vis.begin(), vis.end(), false);

            struct Q {
                Cost key;
                int to;
                bool operator<(Q r) const { return key > r.key; }
            };
            std::priority_queue<Q> que;
            for(int i=0;i<n;i++){
                if(b[i]>=delta){
                    que.push(Q{0, i});
                    dist[i] = 0;
                }
            }

            int t = -1;
            while(!que.empty()){
                int v = que.top().to;
                que.pop();
                if(vis[v]) continue;
                vis[v] = true;
                if(b[v] <= -delta){
                    t = v;
                    break;
                }
                for(int i=0;i<(int)g[v].size();i++){
                    auto &e = g[v][i];
                    if(vis[e.to] or !e.available) continue;
                    Cost cost = e.cost - dual[e.to] + dual[v];
                    if(dist[e.to] - dist[v] > cost){
                        dist[e.to] = dist[v] + cost;
                        pv[e.to] = v;
                        pe[e.to] = i;
                        que.push(Q{dist[e.to], e.to});
                    }
                }
            }

            if(t == -1) return -1;

            for(int v=0;v<n;v++){
                if(!vis[v]) continue;
                dual[v] -= dist[t] - dist[v];
            }

            return t;
        };

        auto delta_scaling_primal = [&](Flow delta, int t){

            int s = -1;
            for(int v=t;b[v]<delta;v=pv[v]){
                auto& e = g[pv[v]][pe[v]];
                auto& e_func = E[e.id];
                auto& re = g[v][e.rev];

                if(e.is_rev) e_func.flow -= delta;
                else e_func.flow += delta;
                delta_update_edge(e, delta);
                delta_update_edge(re, delta);

                if(b[pv[v]] >= delta) s = pv[v];
            }

            assert(s != -1);
            b[s] -= delta;
            b[t] += delta;
        };

        while(delta>0){
            saturate(delta);
            while(true){
                int t = delta_scaling_dual(delta);
                if(t == -1) break;
                delta_scaling_primal(delta, t);
            }
            delta /= 2;
        }

        for(int i=0;i<n;i++) assert(b[i] == 0);

        return;
    }
};

int main(){

    int N, M;
    std::cin >> N >> M;
    std::vector<int> A(N), B(N);
    for(int i=0;i<N;i++) std::cin >> A[i];
    for(int i=0;i<N;i++) std::cin >> B[i];
    std::vector<std::vector<int>> T(N, std::vector<int>(N, 0));
    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++){
            std::cin >> T[i][j];
        }
    }

    // assert
    int A_sum = 0, B_sum = 0;
    for(int i=0;i<N;i++){
        assert(A[i] >= 0);
        A_sum += A[i];
    }
    for(int i=0;i<N;i++){
        assert(B[i] >= 0);
        B_sum += B[i];
    }
    assert(A_sum == M);
    assert(B_sum == M);

    // graph construction
    mccf_graph<int64_t, int64_t> G(2*N);

    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++){
            auto F = [&T, i, j](int64_t x){return (x-T[i][j])*(x-T[i][j]);};
            G.add_edge(i, N+j, F, A[i], 0);
        }
    }

    std::vector<int64_t> C(2*N);
    for(int i=0;i<N;i++){
        C[i] = A[i];
        C[N+i] = -B[i];
    }

    G.solve_CS(C);

    std::vector<std::vector<int>> res(N, std::vector<int>(N, 0));
    for(auto &e: G.E){
        res[e.from][e.to-N] = e.flow;
    }

    // for(int i=0;i<N;i++){
    //     for(int j=0;j<N;j++){
    //         std::cout << res[i][j] << " ";
    //     }
    //     std::cout << "\n";
    // }    

    std::cout << G.calc_objective() << "\n";

    return 0;
}