Untitled

 avatar
unknown
plain_text
a year ago
8.2 kB
5
Indexable
#include <iostream>
#include <vector>
#include <map>
#include <set>

using namespace std;

#define ll long long
#define vi vector<int>
#define pii pair<int, int>
#define vii vector<pii>
#define vli vector<pair<ll, int>>
#define vl vector<long long>
#define pb push_back

const ll INF = 1e18;

vi x, c, y, t;
int n, m, k;
vector <vi> vc;
vi setTo; //вектор для перебора; в какое множество кладем элемент каждого типа
vi numTo;
vector <vli> opt;//для каждой маски храним вектор k наилучших пар (мин суммарное расстояние; цвет)
ll ans;
vii edges;
vector <vi> cap, fl;
vector <vl> cst;
int source, target, all;
int step;

void calc(int mask) {
    vli optc;
    vi cy;
    for (int i = 0; i < m; i++) {
        if (mask & (1 << t[i])) cy.pb(y[i]);
    }
    if (!cy.size()) return;
    for (int clr = 0; clr < n; clr++) {//перебираем цвет
        //решаем задачу распределения шаров цвета clr по лункам типов из mask
        vi cx = vc[clr];
        if (!cx.size()) continue;
        vector<vl> dp(cy.size());
        for (int yy = 0; yy < cy.size(); yy++) {
            dp[yy].resize(cx.size());
            for (int xx = 0; xx < cx.size(); xx++) {
                dp[yy][xx] = INF;
                if (xx) dp[yy][xx] = dp[yy][xx - 1];
                if (xx) {
                    if (yy) dp[yy][xx] = dp[yy - 1][xx - 1] + abs(cy[yy] - cx[xx]);
                    else dp[yy][xx] = abs(cy[yy] - cx[xx]);
                }
                else {
                    dp[yy][xx] = (yy ? INF : abs(cy[yy] - cx[xx]));
                }
            }
        }
        //cout << mask << " " << dp[cy.size() - 1][cx.size() - 1] << " " << clr << endl;
        optc.pb({dp[cy.size() - 1][cx.size() - 1], clr});
    }
    sort(optc.begin(), optc.end());
    for (int i = 0; i < k; i++) opt[mask].pb(optc[i]);
}

pair<int, ll> minCostFlow() {
    pair<int, ll> res = {0, 0ll};
    while (true) {
        vector <ll> d(all);
        for (int i = 0; i < all; i++) d[i] = INF;
        d[source] = 0;
        cout << "all = " << all << endl;
        vector <int> parent(all);
        for (int i = 0; i < all; i++) {
            for (auto [v, u] : edges) {
                if (fl[v][u] == cap[v][u]) continue;
                if (all == 6) cout << d[5] << " " <<  v << " " << u << " " << fl[v][u] << " " << cap[v][u] << " " << cst[v][u] << " " << d[v] << " " << d[u] << endl;
                if (d[v] != INF && d[u] > d[v] + cst[v][u]) {
                    d[u] = d[v] + cst[v][u];
                    parent[u] = v;
                }
            }
        }
        if (d[target] == INF) break;
        cout << "d[target] = " << d[target] << endl;
        int v = target;
        int addFlow = 1e9;
        while (v != source) {
            int u = parent[v];
            addFlow = min(addFlow, cap[u][v] - fl[u][v]);
            v = u;
        }
        v = target;
        res.first += addFlow;
        cout << addFlow << endl;
        while (v != source) {
            int u = parent[v];
            cout << "print = " << u << " " << v << " " << addFlow << " " << cst[u][v] << endl;
            res.second += addFlow * cst[u][v];
            fl[u][v] += addFlow;
            fl[v][u] -= addFlow;
            v = u;
        }
        cout << res.second << endl;
    }
    cout << res.first << " " << res.second << endl;
    return res;
}

ll finalCalc() {
    vi setmasks;
    int sm = 0;
    //cout << "here2" << endl;
    for (int i = 0; i < k; i++) sm = max(sm, setTo[i]);
    //for (int i = 0; i < k; i++) cout << setTo[i] << " ";
    //cout << endl;
    sm++;
    setmasks.resize(sm);
    for (int i = 0; i < k; i++) {
        setmasks[setTo[i]] |= (1 << i);
    }
    //cout << setmasks.size() << " " << setmasks[0] << endl;
    map <int, int> vals;
    set <int> tmp;
    for (int mask : setmasks) {
        for (auto [cost, clr] : opt[mask]) {
            tmp.insert(clr);
        }
    }
    int cc = 0;
    for (auto x : tmp) {
        vals[x] = cc++;
    } 
    all = sm + vals.size() + 2;
    edges.clear();
    for (int i = 0; i < sm; i++) {
        int mask = setmasks[i];
        for (auto [cost, clr] : opt[mask]) {
            cout << mask << " " << cost << " " << clr << " " << vals[clr] << endl;
            edges.pb({i, vals[clr] + sm});
            edges.pb({vals[clr] + sm, i});
            cst[i][vals[clr] + sm] = cost;
            cst[vals[clr] + sm][i] = -cost;
            fl[i][vals[clr] + sm] = fl[vals[clr] + sm][i] = 0;
            cap[i][vals[clr] + sm] = 1;
            cap[vals[clr] + sm][i] = 0;
        }
    }
    source = all - 2;
    target = all - 1;
    cout << "all = " << all << " " << source << " " << target << endl;
    //cout << "here3" << endl;
    for (int i = 0; i < sm; i++) {
        edges.pb({source, i});
        edges.pb({i, source});
        cst[source][i] = cst[i][source] = 0;
        fl[source][i] = fl[i][source] = 0;
        cap[source][i] = 1;
        cap[i][source] = 0;
    }
    //cout << "here4" << endl;
    //cout << sm << " " << vals.size() << endl;
    for (int i = sm; i < sm + vals.size(); i++) {
        edges.pb({i, target});
        edges.pb({target, i});
        //cout << i << " " << target << " " << cst.size() << endl;
        cst[i][target] = cst[target][i] = 0;
        fl[i][target] = fl[target][i] = 0;
        cap[i][target] = 1;
        cap[target][i] = 0;
    }
    for (int i = 0; i < k; i++) cout << setTo[i] << " ";
    cout << endl;
    for (auto [v, u] : edges) cout << v << " " << u << " " << fl[v][u] << " " << cap[v][u] << " " << cst[v][u] << endl;
    cout << endl;
    auto [flow, cost] = minCostFlow();
    if (flow != sm) return INF;
    return cost;
}

//перебор разбиения множества типов на все возможные подмножества
void go(int s, int i, bool was) {
//на каком множестве стоим, на каком типе элементов стоим, непусто ли последнее открытое множество
    cout << "go print = " << s << " " << i << " " << was << endl;
    if (i >= k) {
        if (was) ans = min(ans, finalCalc());
        return;
    }
    for (int is = 0; is <= s; is++) {
        setTo[i] = is;
        go(s, i + 1, was | (is == s));
    }
    if (was) go(s + 1, i, 0);
}

void compress(vector <int> &t) {
    vii p;
    for (int i = 0; i < t.size(); i++) p.pb({t[i], i});
    sort(p.begin(), p.end());
    int val = 0;
    for (int i = 0; i < p.size(); i++) {
        if (i > 0 && p[i].first != p[i - 1].first) val++;
        t[p[i].second] = val;
    }
}

void read() {
    cin >> n >> m >> k;
    x.resize(n);
    c.resize(n);
    y.resize(m);
    t.resize(m);
    for (int i = 0; i < n; i++) cin >> x[i];
    for (int i = 0; i < n; i++) cin >> c[i];
    for (int i = 0; i < n; i++) c[i]--;
    for (int i = 0; i < m; i++) cin >> y[i];
    for (int i = 0; i < m; i++) cin >> t[i];
    for (int i = 0; i < m; i++) t[i]--;
    compress(t);
    k = 0;
    for (int i = 0; i < m; i++) k = max(k, t[i]);
    k++;
}

void solve() {
    read();
    vc.resize(n);
    for (int i = 0; i < n; i++) vc[c[i]].pb(x[i]);
    opt.resize(1 << k);
    for (int tmask = 0; tmask < (1 << k); tmask++) {//перебираем маску типов, которые будут иметь один цвет 
        calc(tmask);
    }
    ans = INF;
    setTo.resize(k);
    int mxsz = k + k * k + 2;
    cap.resize(mxsz);
    fl.resize(mxsz);
    cst.resize(mxsz);
    for (int i = 0; i < mxsz; i++) {
        cap[i].resize(mxsz);
        fl[i].resize(mxsz);
        cst[i].resize(mxsz);
        for (int j = 0; j < mxsz; j++) {
            cap[i][j] = 0;
            fl[i][j] = 0;
            cst[i][j] = 0;
        }
    }
    go(0, 0, 0);
    if (ans == INF) ans = -1;
    cout << ans;
}

int main()
{
    solve();
    return 0;
}
Editor is loading...
Leave a Comment