Untitled

 avatar
unknown
plain_text
2 years ago
6.4 kB
3
Indexable
#include <iostream>

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 << i)) cy.pb(y[i]);
    }
    for (int clr = 0; clr < n; clr++) {//перебираем цвет
        //решаем задачу распределения шаров цвета clr по лункам типов из mask
        vector<vl> dp(m);
        vi cx = vc[clr];
        for (int yy = 0; yy < m; yy++) {
            dp[yy].resize(n);
            for (int xx = 0; xx < n; 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]));
                }
            }
        }
        optc.pb({dp[m - 1][n - 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;
        vector <int> parent(all);
        for (int i = 0; i < n; i++) {
            for (auto [v, u] : edges) {
                if (fl[v][u] == cap[v][u]) continue;
                if (d[u] < d[v] + cst[v][u]) {
                    d[u] = d[v] + cst[v][i];
                    parent[u] = v;
                }
            }
        }
        if (d[target] == INF) break;
        int v = target;
        int addFlow = 0;
        while (v != sink) {
            int u = parent[v];
            addFlow = min(addFlow, cap[u][v] - fl[u][v]);
            v = u;
        }
        v = target;
        res.first += addflow;
        while (v != sink) {
            int u = parent[v];
            res.second += addflow * cst[u][v];
            fl[u][v] += addflow;
            fl[v][u] -= addflow;
            v = u;
        }
    }
    return res;
}

ll finalCalc() {
    vi setmasks;
    int sm = 0;
    for (int i = 0; i < k; i++) sm = max(sm, setTo[i]);
    sm++;
    setmasks.resize(sm);
    for (int i = 0; i < k; i++) {
        setmasks[setTo[i]] |= (1 << i);
    }
    map <int, int> vals;
    vi tmp;
    for (int mask : setmasks) {
        for ([cost, clr] : opt[mask]) {
            tmp.pb(clr);
        }
    }
    sort(tmp.begin(), tmp.end());
    for (int i = 0; i < tmp.size(); i++) vals[tmp[i]] = i;
    all = sm + vals.size() + 2;
    edges.clear();
    for (int i = 0; i < sm; i++) {
        int mask = setmasks[i];
        for (int [cost, clr] : opt[mask]) {
            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] = 0;
            fl[vals[clr] + sm][i] = 0;
            cap[i][vals[clr] + sm] = 1;
            cap[vals[clr] + sm][i] = 0;
        }
    }
    source = all - 2;
    target = all - 1;
    for (int i = 0; i < sm; i++) {
        edges.pb({source, i});
        edges.pb({i, source});
        cst[source][i] = 0;
        cst[i][source] = 0;
        fl[source][i] = 0;
        fl[i][source] = 0;
        cap[source][i] = 1;
        cap[i][source] = 0;
    }
    for (int i = sm; i < sm + vals.size(); i++) {
        edges.pb({i, target});
        edges.pb({target, i});
        cst[i][target] = 0;
        cst[target][i] = 0;
        fl[i][target] = 0;
        fl[target][i] = 0;
        cap[i][target] = 0;
        cap[target][i] = 1;
    }
    auto [flow, cost] = minCostFlow();
    if (flow != sm) return INF;
    return cost;
}

//перебор разбиения множества типов на все возможные подмножества
void go(int s, int i, bool was) {
//на каком множестве стоим, на каком типе элементов стоим, непусто ли последнее открытое множество
    if (i >= k) {
        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 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]--;
}

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);
    cap.resize(k * k + 2);
    fl.resize(k * k + 2);
    cst.resize(k * k + 2);
    marked.resize(k * k + 2);
    for (int i = 0; i < k * k + 2; i++) {
        cap[i].resize(k * k + 2);
        fl[i].resize(k * k + 2);
        cst[i].resize(k * k + 2);
        marked[i].resize(k * k + 2);
        for (int j = 0; j < k * k + 2; 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