Submission #942000

#	Time	Username	Problem	Language	Result	Execution time	Memory
942000		riariti	Toll (APIO13_toll)	C++17	100 / 100	2467 ms	16584 KiB

toll

#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"

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::int32_t main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

#define int std::int64_t

    // 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);

    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:238: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]

• Subtask 1: 16 / 16
• Subtask 2: 18 / 18
• Subtask 3: 22 / 22
• Subtask 4: 22 / 22
• Subtask 5: 22 / 22