Submission #912459

#	Time	Username	Problem	Language	Result	Execution time	Memory
912459		vjudge1	Election Campaign (JOI15_election_campaign)	C++17	100 / 100	325 ms	41956 KiB

election_campaign

// https://oj.uz/problem/view/JOI15_election_campaign

#include <bits/stdc++.h>
using namespace std;

namespace std {

template <int D, typename T>
struct Vec : public vector<Vec<D - 1, T>> {
        static_assert(D >= 1, "Dimension must be positive");
        template <typename... Args>
        Vec(int n = 0, Args... args) : vector<Vec<D - 1, T>>(n, Vec<D - 1, T>(args...)) {}
};

template <typename T>
struct Vec<1, T> : public vector<T> {
        Vec(int n = 0, T val = T()) : std::vector<T>(n, val) {}
};

/* Example
        Vec<4, int64_t> f(n, k, 2, 2); // = f[n][k][2][2];
        Vec<2, int> adj(n); // graph
*/

template <class Fun>
class y_combinator_result {
        Fun fun_;

       public:
        template <class T>
        explicit y_combinator_result(T &&fun) : fun_(std::forward<T>(fun)) {}

        template <class... Args>
        decltype(auto) operator()(Args &&...args) {
                return fun_(std::ref(*this), std::forward<Args>(args)...);
        }
};

template <class Fun>
decltype(auto) y_combinator(Fun &&fun) {
        return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

/* Example
        auto fun = y_combinator([&](auto self, int x) -> void {
                self(x + 1);
        });
*/

}  // namespace std

class segtree_t {
       public:
        segtree_t *left, *right;
        int l, r, m, val, lazy;

        segtree_t(int l, int r) : l(l), r(r), m(l + r >> 1), val(0), lazy(0) {
                if (l == r) return;
                left = new segtree_t(l, m);
                right = new segtree_t(m + 1, r);
        }

        void Update(int s, int t, int x) {
                if (r < s || l > t) return;
                if (s <= l && r <= t) {
                        val += x;
                        lazy += x;
                        return;
                }
                Down();
                left->Update(s, t, x);
                right->Update(s, t, x);
        }

        int Get(int p) {
                if (l == r) return val;
                Down();
                if (p <= m) return left->Get(p);
                return right->Get(p);
        }

        void Down() {
                left->lazy += lazy;
                right->lazy += lazy;
                right->val += lazy;
                left->val += lazy;
                lazy = 0;
        }
};

int32_t main() {
        ios_base::sync_with_stdio(0);
        cin.tie(0);

        int n;
        cin >> n;
        vector<vector<int>> adj(n);
        for (int i = 0; i < n - 1; i++) {
                int u, v;
                cin >> u >> v;
                u--, v--;
                adj[u].emplace_back(v);
                adj[v].emplace_back(u);
        }

        vector<vector<int>> up(__lg(n) + 1, vector<int>(n));
        vector<int> depth(n);
        vector<int> tin(n), tout(n);
        int timer = 0;

        auto dfs = y_combinator([&](auto dfs, int u, int p) -> void {
                tin[u] = timer++;
                up[0][u] = p;
                for (int i = 1; i <= __lg(n); i++) up[i][u] = up[i - 1][up[i - 1][u]];
                for (int v : adj[u]) {
                        if (v == p) continue;
                        depth[v] = depth[u] + 1;
                        dfs(v, u);
                }
                tout[u] = timer;
        });

        dfs(0, 0);

        auto lca = [&](int x, int y) {
                if (depth[x] > depth[y]) swap(x, y);
                for (int i = __lg(n); i >= 0; i--) {
                        if (depth[up[i][y]] >= depth[x]) y = up[i][y];
                }
                if (x == y) return x;
                for (int i = __lg(n); i >= 0; i--) {
                        if (up[i][x] == up[i][y]) continue;
                        x = up[i][x], y = up[i][y];
                }
                return up[0][x];
        };

        int m;
        cin >> m;
        vector<int> x(m), y(m), c(m);
        vector<vector<int>> path(n);

        for (int i = 0; i < m; i++) cin >> x[i] >> y[i] >> c[i], x[i]--, y[i]--;
        for (int i = 0; i < m; i++) path[lca(x[i], y[i])].emplace_back(i);

        segtree_t *tree = new segtree_t(0, n);
        vector<int> dp(n);

        auto dfs_ans = y_combinator([&](auto dfs_ans, int u, int p) -> void {
                int sum = 0;
                for (int v : adj[u]) {
                        if (v == p) continue;
                        dfs_ans(v, u);
                        sum += dp[v];
                        tree->Update(tin[v], tout[v] - 1, -dp[v]);
                }
                tree->Update(tin[u], tin[u], sum);
                dp[u] = sum;
                for (int i : path[u]) {
                        int f = sum + c[i];
                        if (x[i] != u) f += tree->Get(tin[x[i]]);
                        if (y[i] != u) f += tree->Get(tin[y[i]]);
                        dp[u] = max(dp[u], f);
                }
                for (int v : adj[u]) {
                        if (v == p) continue;
                        tree->Update(tin[v], tout[v] - 1, sum);
                }
        });

        dfs_ans(0, 0);
        cout << *max_element(dp.begin(), dp.end());
}

Compilation message (stderr)

election_campaign.cpp: In constructor 'segtree_t::segtree_t(int, int)':
election_campaign.cpp:57:51: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
   57 |         segtree_t(int l, int r) : l(l), r(r), m(l + r >> 1), val(0), lazy(0) {
      |                                                 ~~^~~

• Subtask 1: 10 / 10
• Subtask 2: 5 / 5
• Subtask 3: 5 / 5
• Subtask 4: 30 / 30
• Subtask 5: 10 / 10
• Subtask 6: 40 / 40