Submission #389613

#	Time	Username	Problem	Language	Result	Execution time	Memory
389613		maksim1744	Election Campaign (JOI15_election_campaign)	C++17	100 / 100	348 ms	68932 KiB

election_campaign

/*
    author:  Maksim1744
    created: 14.04.2021 11:51:00
*/

#include "bits/stdc++.h"

using namespace std;

using ll = long long;
using ld = long double;

#define mp   make_pair
#define pb   push_back
#define eb   emplace_back

#define sum(a)     ( accumulate ((a).begin(), (a).end(), 0ll))
#define mine(a)    (*min_element((a).begin(), (a).end()))
#define maxe(a)    (*max_element((a).begin(), (a).end()))
#define mini(a)    ( min_element((a).begin(), (a).end()) - (a).begin())
#define maxi(a)    ( max_element((a).begin(), (a).end()) - (a).begin())
#define lowb(a, x) ( lower_bound((a).begin(), (a).end(), (x)) - (a).begin())
#define uppb(a, x) ( upper_bound((a).begin(), (a).end(), (x)) - (a).begin())

template<typename T>             vector<T>& operator--            (vector<T> &v){for (auto& i : v) --i;            return  v;}
template<typename T>             vector<T>& operator++            (vector<T> &v){for (auto& i : v) ++i;            return  v;}
template<typename T>             istream& operator>>(istream& is,  vector<T> &v){for (auto& i : v) is >> i;        return is;}
template<typename T>             ostream& operator<<(ostream& os,  vector<T>  v){for (auto& i : v) os << i << ' '; return os;}
template<typename T, typename U> pair<T,U>& operator--           (pair<T, U> &p){--p.first; --p.second;            return  p;}
template<typename T, typename U> pair<T,U>& operator++           (pair<T, U> &p){++p.first; ++p.second;            return  p;}
template<typename T, typename U> istream& operator>>(istream& is, pair<T, U> &p){is >> p.first >> p.second;        return is;}
template<typename T, typename U> ostream& operator<<(ostream& os, pair<T, U>  p){os << p.first << ' ' << p.second; return os;}
template<typename T, typename U> pair<T,U> operator-(pair<T,U> a, pair<T,U> b){return mp(a.first-b.first, a.second-b.second);}
template<typename T, typename U> pair<T,U> operator+(pair<T,U> a, pair<T,U> b){return mp(a.first+b.first, a.second+b.second);}
template<typename T, typename U> void umin(T& a, U b){if (a > b) a = b;}
template<typename T, typename U> void umax(T& a, U b){if (a < b) a = b;}

#ifdef HOME
#define SHOW_COLORS
#include "C:/C++ libs/print.cpp"
#else
#define show(...)     42
#define mclock        42
#define shows         42
#define debug if (false)
#endif

vector<int> lca_ind;
vector<vector<int>> lca_sparse;
vector<int> lca_p2;
vector<int> lca_depth;
void build_lca_sparse(vector<vector<int>>& g, int root = 0) {
    int n = g.size();
    vector<int> euler;
    lca_ind.resize(n);
    lca_depth.assign(n, -1);
    function<void(int, int)> dfs = [&](int v, int depth) {
        lca_ind[v] = euler.size();
        euler.pb(v);
        lca_depth[v] = depth;
        for (auto k : g[v]) {
            if (lca_depth[k] == -1) {
                dfs(k, depth + 1);
                euler.pb(v);
            }
        }
    };
    dfs(root, 0);
    int m = euler.size();
    int log = 1;
    while ((1 << log) < m)
        ++log;
    lca_sparse.resize(log);
    lca_sparse[0].resize(m);
    lca_p2.resize(m + 1);
    int pp2 = 0;
    for (int i = 1; i < lca_p2.size(); ++i) {
        if (1 << (pp2 + 1) <= i)
            ++pp2;
        lca_p2[i] = pp2;
    }
    lca_p2[0] = 0;
    for (int i = 0; i < m; ++i)
        lca_sparse[0][i] = euler[i];
    for (int i = 1; i < log; ++i) {
        lca_sparse[i].assign(m, 0);
        for (int j = 0; j < m - (1 << (i - 1)); ++j) {
            int v1 = lca_sparse[i - 1][j], v2 = lca_sparse[i - 1][j + (1 << (i - 1))];
            if (lca_depth[v1] < lca_depth[v2])
                lca_sparse[i][j] = v1;
            else
                lca_sparse[i][j] = v2;
        }
    }
}

int get_lca(int u, int v) {
    if (u == v)
        return u;
    u = lca_ind[u];
    v = lca_ind[v];
    if (u > v)
        swap(u, v);
    int v1 = lca_sparse[lca_p2[v - u + 1]][u], v2 = lca_sparse[lca_p2[v - u + 1]][v - (1 << lca_p2[v - u + 1]) + 1];
    if (lca_depth[v1] < lca_depth[v2])
        return v1;
    else
        return v2;
}

int dist(int u, int v) {
    return lca_depth[u] + lca_depth[v] - 2 * lca_depth[get_lca(u, v)];
}

struct item {
    int sm = 0;
    int md = 0;

    template<typename T>
    void init(const T &t, int l, int r) {
        sm = t;
        md = 0;
    }

    void update(const item &first, const item &second, int l, int r) {
        sm = first.sm + second.sm;
    }

    static item merge(const item &first, const item &second, int l, int r) {
        item res;
        res.update(first, second, l, r);  // careful with different lengths
        return res;
    }

    template<typename Modifier>
    void modify(const Modifier &m, int l, int r) {
        // apply here, save for children
        md += m;
        sm += m * (r - l + 1);
    }

    void push(item &first, item &second, int l, int r) {
        int m = (l + r) / 2;
        first.modify(md, l, m);
        second.modify(md, m + 1, r);
        // reset modifier
        md = 0;
    }
};

string to_string(const item &i) {
    stringstream ss;
    ss << "[" << "]";
    return ss.str();
}
ostream& operator << (ostream &o, const item &i) {
    return o << to_string(i);
}

struct segtree {
    vector<item> tree;
    int n = 1;

    segtree(int n = 1) : n(n) {
        tree.resize(1 << (__lg(n - 1) + 2));
    }

    template<typename T>
    void build(const vector<T> &v, int i, int l, int r) {
        if (l == r) {
            tree[i].init(v[l], l, r);
            return;
        }
        int m = (l + r) >> 1;
        build(v, i * 2 + 1, l, m);
        build(v, i * 2 + 2, m + 1, r);
        tree[i].update(tree[i * 2 + 1], tree[i * 2 + 2], l, r);
    }

    template<typename T>
    void build(const vector<T> &v) {
        n = v.size();
        tree.resize(1 << (__lg(n - 1) + 2));
        build(v, 0, 0, n - 1);
    }

    item ask(int l, int r, int i, int vl, int vr) {
        if (vl != vr) {
            tree[i].push(tree[i * 2 + 1], tree[i * 2 + 2], vl, vr);
        }
        if (l == vl && r == vr) {
            return tree[i];
        }
        int m = (vl + vr) >> 1;
        if (r <= m) {
            return ask(l, r, i * 2 + 1, vl, m);
        } else if (m < l) {
            return ask(l, r, i * 2 + 2, m + 1, vr);
        } else {
            return item::merge(ask(l, m, i * 2 + 1, vl, m), ask(m + 1, r, i * 2 + 2, m + 1, vr), l, r);
        }
    }

    item ask(int l, int r) {
        l = max(l, 0); r = min(r, n - 1);
        if (l > r) return item();
        return ask(l, r, 0, 0, n - 1);
    }

    template<typename T>
    void set(int ind, const T &t) {
        static array<pair<int, int>, 30> st;
        int l = 0, r = n - 1, i = 0;
        int ptr = -1;
        while (l != r) {
            if (l != r) {
                tree[i].push(tree[i * 2 + 1], tree[i * 2 + 2], l, r);
            }
            st[++ptr] = {l, r};
            int m = (l + r) >> 1;
            if (ind <= m) {
                i = i * 2 + 1;
                r = m;
            } else {
                i = i * 2 + 2;
                l = m + 1;
            }
        }
        tree[i].init(t, l, r);
        while (i != 0) {
            i = (i - 1) / 2;
            tree[i].update(tree[i * 2 + 1], tree[i * 2 + 2], st[ptr].first, st[ptr].second);
            --ptr;
        }
    }

    template<typename Modifier>
    void modify(int l, int r, const Modifier &modifier, int i, int vl, int vr) {
        if (vl != vr) {
            tree[i].push(tree[i * 2 + 1], tree[i * 2 + 2], vl, vr);
        }
        if (l == vl && r == vr) {
            tree[i].modify(modifier, vl, vr);
            return;
        }
        int m = (vl + vr) >> 1;
        if (r <= m) {
            modify(l, r, modifier, i * 2 + 1, vl, m);
        } else if (m < l) {
            modify(l, r, modifier, i * 2 + 2, m + 1, vr);
        } else {
            modify(l, m, modifier, i * 2 + 1, vl, m);
            modify(m + 1, r, modifier, i * 2 + 2, m + 1, vr);
        }
        tree[i].update(tree[i * 2 + 1], tree[i * 2 + 2], vl, vr);
    }

    template<typename Modifier>
    void modify(int l, int r, const Modifier &modifier) {
        l = max(l, 0); r = min(r, n - 1);
        if (l > r) return;
        modify(l, r, modifier, 0, 0, n - 1);
    }
};

int main() {
    ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);

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

    build_lca_sparse(g, 0);

    int m;
    cin >> m;

    vector<vector<pair<pair<int, int>, int>>> here(n);

    for (int i = 0; i < m; ++i) {
        int a, b, c;
        cin >> a >> b >> c;
        --a; --b;
        here[get_lca(a, b)].eb(mp(a, b), c);
    }

    vector<int> dp(n, 0);
    vector<bool> u(n, false);
    // ind is index in euler tour, [L[v], R[v]] is a segment in this tour corresponding to subtree of vertex v
    vector<int> L(n), R(n), ind(n);
    vector<int> lvl(n, 0);

    segtree tree(n);

    int icur = 0;

    // binary lifting to be able to find child of v, which has vertex x in a subtree
    vector<vector<int>> up(20, vector<int>(n, -1));
    {
        vector<bool> u(n, false);
        function<void(int)> dfs1 = [&](int v) {
            u[v] = true;
            for (int k : g[v]) {
                if (!u[k]) {
                    up[0][k] = v;
                    dfs1(k);
                }
            }
        };
        dfs1(0);
        for (int i = 1; i < up.size(); ++i)
            for (int j = 0; j < n; ++j)
                if (up[i - 1][j] != -1)
                    up[i][j] = up[i - 1][up[i - 1][j]];
    }
    auto go_up = [&](int v, int k) {
        for (int i = 0; i < up.size(); ++i)
            if ((k >> i) & 1)
                v = up[i][v];
        return v;
    };

    function<void(int)> dfs = [&](int v) {
        u[v] = true;
        ind[v] = icur;
        ++icur;
        L[v] = R[v] = ind[v];
        int dp0 = 0;
        vector<int> ch;
        for (int k : g[v]) {
            if (!u[k]) {
                ch.pb(k);
                lvl[k] = lvl[v] + 1;
                dfs(k);
                L[v] = min(L[v], L[k]);
                R[v] = max(R[v], R[k]);
                dp0 += dp[k];
            }
        }
        dp[v] = dp0;
        for (auto [ab, c] : here[v]) {
            auto [a, b] = ab;
            int cur = dp0;
            for (int ch : vector<int>{a, b}) {
                if (ch != v) {
                    // subtract answer for a child of v, containing ch in a subtree
                    cur -= dp[go_up(ch, lvl[ch] - lvl[v] - 1)];
                    // add value along the path
                    cur += tree.ask(ind[ch], ind[ch]).sm;
                }
            }
            dp[v] = max(dp[v], cur + c);
        }
        for (int k : ch) {
            tree.modify(L[k], R[k], dp0 - dp[k]);
        }
        tree.set(ind[v], dp0);
    };

    dfs(0);
    cout << dp[0] << '\n';

    return 0;
}

Compilation message (stderr)

election_campaign.cpp: In function 'void build_lca_sparse(std::vector<std::vector<int> >&, int)':
election_campaign.cpp:77:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   77 |     for (int i = 1; i < lca_p2.size(); ++i) {
      |                     ~~^~~~~~~~~~~~~~~
election_campaign.cpp: In function 'int main()':
election_campaign.cpp:318:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  318 |         for (int i = 1; i < up.size(); ++i)
      |                         ~~^~~~~~~~~~~
election_campaign.cpp: In lambda function:
election_campaign.cpp:324:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  324 |         for (int i = 0; i < up.size(); ++i)
      |                         ~~^~~~~~~~~~~

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