Submission #941999

#TimeUsernameProblemLanguageResultExecution timeMemory
941999riaritiToll (APIO13_toll)C++17
100 / 100
2308 ms24896 KiB
#line 1 "Toll.cpp"
#include <bits/stdc++.h>

#line 2 "/Users/jeff/Documents/git_repos/competitive/cplib/environment/dbg.hpp"

#if defined(local)
#include <debug>
#else
#define dbg(...)
#endif
#line 4 "Toll.cpp"

#define int std::int64_t

using namespace std;
typedef long long lint;
typedef pair<int, int> pi;

struct disj {
    int pa[100005];
    void init(int n) {
        for (int i = 1; i <= n; i++) {
            pa[i] = i;
        }
    }
    int find(int x) { return pa[x] = (pa[x] == x ? x : find(pa[x])); }
    bool uni(int p, int q) {
        p = find(p);
        q = find(q);
        if (p == q)
            return 0;
        pa[q] = p;
        return 1;
    }
} disj;

struct edg {
    int s, e, x;
    bool operator<(const edg &e) const { return x < e.x; }
} a[300005];

int n, m, k, s[22], e[22], c[100005];
vector<int> gph[100005];

vector<edg> resi;
int p, comp[100005];
lint csize[25];

int par[25], pae[25], chk[25], dep[25];
lint sz[25];
vector<pi> gph2[25];

void dfs2(int x, int p) {
    sz[x] = csize[x];
    for (auto &i : gph2[x]) {
        if (i.first == p)
            continue;
        chk[i.first] = i.second;
        dep[i.first] = dep[x] + 1;
        pae[i.first] = 1e9;
        par[i.first] = x;
        dfs2(i.first, x);
        sz[x] += sz[i.first];
    }
}

lint solve(int b) {
    dbg("solve");
    for (int i = 1; i <= p; i++) {
        gph2[i].clear();
    }
    disj.init(p);
    for (int i = 0; i < k; i++) {
        if ((b >> i) & 1) {
            if (!disj.uni(s[i], e[i])) {
                return -1;
            }

            dbg(s[i], e[i]);

            gph2[s[i]].push_back({e[i], 1});
            gph2[e[i]].push_back({s[i], 1});
        }
    }
    vector<edg> resi2;
    for (auto &i : resi) {
        if (disj.uni(i.s, i.e)) {
            dbg(i.s, i.e);

            gph2[i.s].push_back({i.e, 0});
            gph2[i.e].push_back({i.s, 0});
        } else {
            dbg("resi2", i.s, i.e);
            resi2.push_back(i);
        }
    }
    dfs2(1, -1);

    for (int i = 1; i <= n; i++) {
        dbg(i);
        dbg(par[i]);
        dbg(chk[i]);
        dbg(dep[i]);
        dbg(sz[i]);
        dbg(csize[i]);
        dbg(pae[i]);
    }

    for (auto &i : resi2) {
        if (dep[i.s] < dep[i.e])
            swap(i.s, i.e);
        while (dep[i.s] > dep[i.e]) {
            pae[i.s] = min(pae[i.s], i.x);
            i.s = par[i.s];
        }
        while (i.s != i.e) {
            pae[i.s] = min(pae[i.s], i.x);
            pae[i.e] = min(pae[i.e], i.x);
            i.s = par[i.s];
            i.e = par[i.e];
        }
    }
    lint ret = 0;
    for (int i = 2; i <= p; i++) {
        if (chk[i])
            ret += sz[i] * pae[i];
    }
    return ret;
}

void dfs(int x, int p) {
    if (comp[x])
        return;
    csize[p] += c[x];
    comp[x] = p;
    for (auto &i : gph[x]) {
        dfs(i, p);
    }
}

void compress() {
    disj.init(n);
    for (int i = 0; i < k; i++) {
        dbg(s[i], e[i]);
        disj.uni(s[i], e[i]);
    }
    for (int i = 0; i < n - 1; i++) {
        if (disj.uni(a[i].s, a[i].e)) {
            dbg(a[i].e, a[i].s);

            gph[a[i].e].push_back(a[i].s);
            gph[a[i].s].push_back(a[i].e);
        } else {
            dbg("resi", a[i].e, a[i].s);

            resi.push_back(a[i]);
        }
    }

    for (int i = 1; i <= n; i++) {
        if (!comp[i]) {
            dfs(i, ++p);
        }
    }

    for (int i = 1; i <= n; i++) {
        dbg(comp[i]);
    }

    for (auto &i : resi) {
        i.s = comp[i.s];
        i.e = comp[i.e];

        dbg(i.s, i.e);
    }
    for (int i = 0; i < k; i++) {
        s[i] = comp[s[i]];
        e[i] = comp[e[i]];

        dbg(s[i], e[i]);
    }
}

void sol(int N, int M, int K, std::vector<std::array<int, 3>> &abc,
         std::vector<std::array<int, 3>> &xy, std::vector<int> &C) {
    n = N;
    m = M;
    k = K;
    dbg(n, m, k);
    for (int i = 0; i < m; i++) {
        a[i].s = abc[i][0];
        a[i].e = abc[i][1];
        a[i].x = abc[i][2];
        dbg(a[i].s, a[i].e, a[i].x);
    }
    for (int i = 0; i < k; i++) {
        s[i] = xy[i][0];
        e[i] = xy[i][1];
        dbg(s[i], e[i]);
    }
    for (int i = 1; i <= n; i++) {
        c[i] = C[i - 1];
    }
    sort(a, a + m);
    disj.init(n);

    dbg("here");
    vector<edg> tree;
    dbg("here");
    int cnt = 0;
    for (int i = 0; i < m; i++) {
        dbg("hi");
        if (disj.uni(a[i].s, a[i].e)) {
            dbg("hi", a[i].s, a[i].e);
            cnt++;
            tree.push_back(a[i]);
        }
    }
    dbg("here", cnt);
    for (int i = 0; i < n - 1; i++) {
        a[i] = tree[i];
    }

    dbg("here");

    compress();

    lint ret = 0;
    for (int i = 1; i < (1 << k); i++) {
        ret = max(ret, solve(i));
    }
    dbg(ret);

    dbg("-----------------------");
}

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

    // https://dmoj.ca/problem/apio13p2/editorial
    // getting used to style :()
    // got until observation 3 without any doubts

    // // kruskal-ish way
    // std::sort(abc.begin(), abc.end(),
    //           [&](auto x, auto y) { return x[2] < y[2]; });

    // // K <= 20 => 2^K = 1048576
    // // bitmasks which of them we should take then put max value so that they
    // are
    // // still included in minimum spanning tree

    // const auto k = 1 << K;

    // // store the max weight so that edge is contained only compared to the
    // // other constant it possibly does not change what about cases where it
    // is
    // // compared to other non-constants
    //    // kruskal on original edges
    // connect 2 components (see if there were additional edges)
    // atrribute that you can replace the edge with the additional edges (now
    // are a group of edges)
    // choose any 1 from the group
    // weight = to original weight
    // nah

    constexpr int INF = 1E9 + 9;

    int N, M, K;
    std::cin >> N >> M >> K;
    dbg(N, M, K);
    std::vector<std::array<int, 3>> abc(M);
    for (int i = 0; i < M; i++) {
        auto &[a, b, c] = abc[i];
        std::cin >> a >> b >> c;
        // a--;
        // b--;
    }
    std::vector<std::array<int, 3>> xy(K);
    for (int i = 0; i < K; i++) {
        auto &[x, y, w] = xy[i];
        std::cin >> x >> y;
        // w = INF;
        // x--;
        // y--;
    }
    std::vector<int> c(N);
    for (int i = 0; i < N; i++) {
        std::cin >> c[i];
    }

    // sol(N, M, K, abc, xy, c);

    for (int i = 0; i < M; i++) {
        auto &[a, b, c] = abc[i];
        a--;
        b--;
    }
    for (int i = 0; i < K; i++) {
        auto &[x, y, w] = xy[i];

        w = INF;
        x--;
        y--;
    }
    dbg(abc);
    dbg(xy);
    dbg(c);

    struct DSU {
        int n;
        std::vector<int> f;
        std::vector<int> siz;

        DSU() {}
        explicit DSU(int _n) : n(_n), f(_n), siz(_n) {
            std::iota(f.begin(), f.end(), 0);
            siz.assign(n, -1);
        }

        int find(int x) {
            while (x != f[x]) {
                x = f[x] = f[f[x]];
            }

            return x;
        }

        bool same(int x, int y) { return find(x) == find(y); }

        bool merge(int x, int y) {
            x = find(x);
            y = find(y);

            if (x == y) {
                return false;
            }

            siz[x] += siz[y];
            f[y] = x;

            return true;
        }

        int size(int x) { return siz[find(x)]; }
    };

    std::vector<std::array<int, 3>> mst;
    auto kruskal = [&]() -> void {
        // make mst
        std::sort(abc.begin(), abc.end(),
                  [](const auto &x, const auto &y) { return x[2] < y[2]; });
        dbg(abc);

        // kruskal
        DSU dsu(N);
        for (int i = 0; i < M; i++) {
            auto [a, b, c] = abc[i];

            if (dsu.merge(a, b)) {
                mst.push_back({a, b, c});
            }
        }
        dbg(mst);
    };
    kruskal();

    // dont need the bad ori edges / extra useless weight
    abc = std::move(mst);
    mst.clear();

    dbg(abc);

    // M = N - 1
    M = abc.size();
    dbg(M);

    // compress into at max k components (since that is what the questin says)
    // => local (inside) and global path for each component
    // local does not change, at max k vertexs, one from each from current to
    // outside component (at max k components)
    // value is from root until the local root (root of the compoent ) is called
    // global path else inside it is called local path (computation to all can
    // be done via O(n) + O(1) of lca or O(nk) with simple dfs or O(n) with dfs
    // + rerooting)
    // local cost does not depend on query but global does => lcoal cost precomp
    // + global cost recalc with dfs O(k) vertexs only can be done in O(nk + 2^k
    // * k) or O(n + 2^k * k) both passable

    std::vector<std::vector<int>> adj(N);
    // components mapping, number of people
    std::vector<int> cmp(N, -1), np(N);
    int C = -1;
    // not part of k vertexs
    std::vector<std::array<int, 3>> npk;
    // compression to differniate the k edges and original edges and then group
    // them up into components for quicker query since local does not change
    // only global does, and so we have to compress as mentioned above
    auto compress = [&]() -> void {
        DSU dsu(N);

        // choose the vertexs which get affected by the k new nodes
        // / the k edges themselves
        for (int i = 0; i < K; i++) {
            auto [x, y, w] = xy[i];
            dbg(x, y, w);

            if (dsu.merge(x, y)) {
                dbg(x, y);
            }
        }

        // go through the mst to see which paths between vertexs get affected
        for (int i = 0; i < M; i++) {
            auto [a, b, c] = abc[i];

            if (dsu.merge(a, b)) {
                dbg(a, b);
                // ok so not affected => take
                // dont put the graph edge of x y
                // since we want to find the compoents where it is seperated
                adj[a].push_back(b);
                adj[b].push_back(a);
            } else {
                // affeceted => cant take => but still compress
                // to see where the nodes end up at i.e which component
                // part of the components centered by the above k vertexs

                // will be used especially since we are not going to take all of
                // the K all the K edges, sometimes only a subset
                npk.push_back({a, b, c});
            }
        }
        dbg(npk);

        // fill in components i.e what component will vertex a, b and so on
        // belong to and how many people are from there
        dbg(cmp);
        for (int i = 0, p = 0; i < N; i++) {
            if (cmp[i] != -1) {
                continue;
            }

            auto dfs = [&](auto &&self, int x) {
                // baiscally like vis
                if (cmp[x] != -1) {
                    return;
                }

                dbg(np[p], c[x]);
                np[p] += c[x];
                cmp[x] = p;

                for (auto y : adj[x]) {
                    self(self, y);
                }
            };
            dfs(dfs, i);

            p++;
            C = p;
        }
        dbg(cmp, C);
        dbg(np);

        // convert them to components
        for (int i = 0; i < npk.size(); i++) {
            auto &[a, b, c] = npk[i];

            a = cmp[a];
            b = cmp[b];

            dbg(a, b);
        }

        // convert the k edges into which compoentes they belng to
        for (int i = 0; i < K; i++) {
            auto &[x, y, w] = xy[i];

            x = cmp[x];
            y = cmp[y];

            dbg(x, y);
        }
    };
    compress();

    auto sol = [&](const std::vector<int> &idx) -> std::int64_t {
        dbg(idx);

        DSU dsu(C);
        std::vector<std::vector<std::pair<int, int>>> gdj(C);

        for (auto i : idx) {
            auto [x, y, w] = xy[i];

            // it should not be a cycle
            if (!dsu.merge(x, y)) {
                return -1;
            }

            dbg(x, y);

            // 1 weight and 0 for the non (chosen) k edges and original edges
            // for check function since we only can add values if they pass
            // through this edge
            gdj[x].push_back({y, 1});
            gdj[y].push_back({x, 1});
        }
        // dbg(gdj);

        // again similar with compression, but now only with the non chosen k
        // edges
        std::vector<std::array<int, 3>> npl;
        // make sure onyl from the original non (all) k edges
        for (auto [a, b, c] : npk) {
            dbg(a, b, c);

            if (dsu.merge(a, b)) {
                dbg("here", a, b);
                // 0 for the non profit edges
                gdj[a].push_back({b, 0});
                gdj[b].push_back({a, 0});
            } else {
                npl.push_back({a, b, c});
            }
        }
        dbg(npl);

        std::vector<int> par(C, -1), chk(C), dep(C), siz(C);
        // np = weight of edges that will get paid from
        // pae = all k's edges or nah (will be used for cost calc or nah)
        std::vector<int> pae(C);
        // std::vector<int> np(N), pae(N);
        // np alreadyinitialized
        // O(K) vertexs
        auto dfs = [&](auto &&self, int x) -> void {
            // siz is number fo people
            siz[x] = np[x];
            for (auto [y, w] : gdj[x]) {
                if (y == par[x]) {
                    continue;
                }

                par[y] = x;
                chk[y] = w;
                dep[y] = dep[x] + 1;
                pae[y] = INF;

                self(self, y);

                siz[x] += siz[y];
            }
        };
        dfs(dfs, 0);

        dbg(par);
        dbg(chk);
        dbg(dep);
        dbg(siz);
        dbg(np);
        dbg(pae);

        for (auto [a, b, c] : npl) {
            // mins here are to check if all of them are have weight
            // the min weight here is the dependent on the k edges
            // so the win weight / cost of the k edges here are the 0's and 1's
            // basically checking if it will benefit him or nah

            // if it passess through only the k edgees then only will the val
            // will be 1 else 0 i.e the person wont get any money
            // and it si possible since siz is basicaly like a prefix sum
            // so if we add it onyl where it passes through the kedges => it is
            // still answer since it tells how many people pass through it
            // confirm

            // lca path checking basically
            if (dep[a] < dep[b]) {
                std::swap(a, b);
            }

            // equalize depth
            while (dep[a] > dep[b]) {
                pae[a] = std::min(pae[a], c);
                a = par[a];
            }

            // lca
            while (a != b) {
                pae[a] = std::min(pae[a], c);
                pae[b] = std::min(pae[b], c);

                a = par[a];
                b = par[b];
            }
        }

        int ans = 0;
        for (int i = 1; i < C; i++) {
            if (!chk[i]) {
                continue;
            }

            ans += siz[i] * pae[i];
        }

        return ans;
    };

    const auto bit_set = [](int v, int k) -> bool { return v & (1 << k); };

    std::int64_t ans = 0;
    for (int k = 0; k < (1 << K); k++) {
        std::vector<int> idx;
        for (int i = 0; i < K; i++) {
            if (!bit_set(k, i)) {
                continue;
            }

            idx.push_back(i);
        }

        ans = std::max(ans, sol(idx));
    }

    std::cout << ans << "\n";

#undef int

    return 0;
}

Compilation message (stderr)

Toll.cpp: In lambda function:
Toll.cpp:460:27: warning: comparison of integer expressions of different signedness: 'int64_t' {aka 'long int'} and 'std::vector<std::array<long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...