Submission #1116885

#	Time	Username	Problem	Language	Result	Execution time	Memory
1116885		Zero_OP	Election Campaign (JOI15_election_campaign)	C++14	100 / 100	130 ms	54468 KiB

election_campaign

#include <bits/stdc++.h>

using namespace std;

#define dbg(x) "[" #x " = " << (x) << "]"

template<typename T>
bool maximize(T& a, const T& b){
    if(a < b) {
        return a = b, true;
    } return false;
}

struct fenwick_tree{
    vector<int> bit;
    fenwick_tree() : bit() {}
    fenwick_tree(int n) : bit(n + 1, 0) {}

    void update(int i, int v){
        for(; i < (int)bit.size(); i += i & (-i)) bit[i] += v;
    }

    void update(int l, int r, int v){
        update(l, +v);
        update(r + 1, -v);
    }

    int query(int i){
        int sum = 0;
        for(; i > 0; i -= i & (-i)) sum += bit[i];
        return sum;
    }

    int query(int l, int r){
        return query(r) - query(l - 1);
    }
};

struct heavy_light_decomposition{
    int N, timerHLD;
    vector<int> depth, par, head, sz, tin, tout, dp, sum_dp;
    vector<vector<int>> adj;
    vector<vector<pair<int, int>>> one_subtree;
    vector<vector<tuple<int, int, int>>> two_subtrees;
    fenwick_tree ft;

    heavy_light_decomposition(int N) :
        timerHLD(0),
        N(N),
        depth(N),
        par(N, -1),
        head(N, -1),
        sz(N, -1),
        adj(N),
        tin(N),
        tout(N),
        one_subtree(N),
        two_subtrees(N),
        dp(N, 0),
        sum_dp(N, 0),
        ft(N) {}

    void add_edge(int u, int v){
        adj[u].emplace_back(v);
        adj[v].emplace_back(u);
//        cout << dbg(u) << dbg(v) << '\n';
    }

    void dfs_sz(int u){
        sz[u] = 1;
        for(auto& v : adj[u]){
            adj[v].erase(find(adj[v].begin(), adj[v].end(), u));
            depth[v] = depth[u] + 1;
            par[v] = u;
            dfs_sz(v);
            sz[u] += sz[v];
            if(sz[v] > sz[adj[u][0]]) swap(v, adj[u][0]);
        }
    }

    void dfs_hld(int u, int hd){
        tin[u] = ++timerHLD;
        head[u] = hd;

        for(auto v : adj[u]){
            if(v == adj[u][0]) dfs_hld(v, hd);
            else dfs_hld(v, v);
        }

        tout[u] = timerHLD;
    }

    void preprocess(int rt = 0){
        dfs_sz(rt);
        dfs_hld(rt, rt);
    }

    bool in_subtree(int u, int v){
        return tin[u] <= tin[v] && tout[v] <= tout[u];
    }

    int get_lca(int u, int v){
        if(in_subtree(u, v)) return u;
        if(in_subtree(v, u)) return v;

        while(head[u] != head[v]){
            if(depth[head[u]] < depth[head[v]]) swap(u, v);
            u = par[head[u]];
        }

        if(tin[u] > tin[v]) swap(u, v);
        return u;
    }

    void add_path(int u, int v, int w){
        if(tin[u] > tin[v]) swap(u, v);
        int lca = get_lca(u, v);

        if(lca == u){
            one_subtree[lca].emplace_back(v, w);
        } else{
            two_subtrees[lca].emplace_back(u, v, w);
        }
    }

    int find_exact_subtree(int parent, int u){
        int l = 0, r = (int)adj[parent].size() - 1, ans = -1;
        while(l <= r){
            int mid = l + r >> 1;
            if(tin[adj[parent][mid]] <= tin[u]) ans = mid, l = mid + 1;
            else r = mid - 1;
        }

        assert(ans != -1);
        return adj[parent][ans];
    }

    void update(int u, int val){
        ft.update(tin[u], val);
    }

    int sum_path(int parent, int u){
        if(depth[u] - depth[parent] < 0) return 0;
        int res = 0;
        while(head[parent] != head[u]){
            res += ft.query(tin[head[u]], tin[u]);
            u = par[head[u]];
        }

        assert(tin[parent] <= tin[u]);
        res += ft.query(tin[parent], tin[u]);
        return res;
    }

    void solve(int u){
        sum_dp[u] = 0;
        for(int v : adj[u]){
            solve(v);
            sum_dp[u] += dp[v];
        }

        //case 1 : skip u
        maximize(dp[u], sum_dp[u]);

        //case 2 : consider every paths where lca(A, B) = u
        ///case 2.1 : A = u or B = u
        for(auto [x, w] : one_subtree[u]){
            int px = find_exact_subtree(u, x);

            int cur = sum_dp[u] + sum_dp[x] + sum_path(px, x) - dp[px];
            maximize(dp[u], cur + w);
        }

        ///case 2.2 : A \neq u and B \neq u
        for(auto [x, y, w] : two_subtrees[u]){
            int px = find_exact_subtree(u, x);
            int py = find_exact_subtree(u, y);

            int cur = sum_dp[u] + sum_dp[x] + sum_dp[y] - dp[px] - dp[py] + sum_path(px, x) + sum_path(py, y);
            maximize(dp[u], cur + w);
        }

        for(auto v : adj[u]){
            update(v, sum_dp[u] - dp[v]);
        }
    }
};

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

#ifdef LOCAL
    freopen("task.inp", "r", stdin);
    freopen("task.out", "w", stdout);
#endif // LOCAL

    int N; cin >> N;
    heavy_light_decomposition tr(N);
    for(int i = 1; i < N; ++i){
        int u, v;
        cin >> u >> v;
        --u, --v;
        tr.add_edge(u, v);
    }

    tr.preprocess();

    int M; cin >> M;
    vector<tuple<int, int, int>> paths;
    for(int i = 0; i < M; ++i){
        int u, v, w;
        cin >> u >> v >> w;
        --u, --v;
        tr.add_path(u, v, w);
    }

    tr.solve(0);
    cout << tr.dp[0] << '\n';

    return 0;
}

Compilation message (stderr)

election_campaign.cpp: In constructor 'heavy_light_decomposition::heavy_light_decomposition(int)':
election_campaign.cpp:40:12: warning: 'heavy_light_decomposition::timerHLD' will be initialized after [-Wreorder]
   40 |     int N, timerHLD;
      |            ^~~~~~~~
election_campaign.cpp:40:9: warning:   'int heavy_light_decomposition::N' [-Wreorder]
   40 |     int N, timerHLD;
      |         ^
election_campaign.cpp:47:5: warning:   when initialized here [-Wreorder]
   47 |     heavy_light_decomposition(int N) :
      |     ^~~~~~~~~~~~~~~~~~~~~~~~~
election_campaign.cpp:42:25: warning: 'heavy_light_decomposition::adj' will be initialized after [-Wreorder]
   42 |     vector<vector<int>> adj;
      |                         ^~~
election_campaign.cpp:41:39: warning:   'std::vector<int> heavy_light_decomposition::tin' [-Wreorder]
   41 |     vector<int> depth, par, head, sz, tin, tout, dp, sum_dp;
      |                                       ^~~
election_campaign.cpp:47:5: warning:   when initialized here [-Wreorder]
   47 |     heavy_light_decomposition(int N) :
      |     ^~~~~~~~~~~~~~~~~~~~~~~~~
election_campaign.cpp:44:42: warning: 'heavy_light_decomposition::two_subtrees' will be initialized after [-Wreorder]
   44 |     vector<vector<tuple<int, int, int>>> two_subtrees;
      |                                          ^~~~~~~~~~~~
election_campaign.cpp:41:50: warning:   'std::vector<int> heavy_light_decomposition::dp' [-Wreorder]
   41 |     vector<int> depth, par, head, sz, tin, tout, dp, sum_dp;
      |                                                  ^~
election_campaign.cpp:47:5: warning:   when initialized here [-Wreorder]
   47 |     heavy_light_decomposition(int N) :
      |     ^~~~~~~~~~~~~~~~~~~~~~~~~
election_campaign.cpp: In member function 'int heavy_light_decomposition::find_exact_subtree(int, int)':
election_campaign.cpp:129:25: warning: suggest parentheses around '+' inside '>>' [-Wparentheses]
  129 |             int mid = l + r >> 1;
      |                       ~~^~~
election_campaign.cpp: In member function 'void heavy_light_decomposition::solve(int)':
election_campaign.cpp:167:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  167 |         for(auto [x, w] : one_subtree[u]){
      |                  ^
election_campaign.cpp:175:18: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  175 |         for(auto [x, y, w] : two_subtrees[u]){
      |                  ^

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