Submission #130832

#	Time	Username	Problem	Language	Result	Execution time	Memory
130832		xantho	Election Campaign (JOI15_election_campaign)	C++17	100 / 100	330 ms	31528 KiB

election_campaign

#include <algorithm>
#include <iostream>
#include <vector>

class heavy_light {
    struct path {
        int u, v;
        int lca;
        long long value;
    };

    enum class chain_value_type {
        children,
        vertex
    };

    const int UNCHECKED = -1;
    const int ROOT = 0; // Standard root

    int num_vertices;
    std::vector<std::vector<int>> adj_list;
    std::vector<int> parent; // Parent of each vertex in rooted tree.
    std::vector<int> depth; // Depth of each vertex.
    std::vector<int> dfs_start; // Time when a vertex's DFS processing starts. Used to check for Ancestors
    std::vector<int> dfs_end; // Time when a vertex's DFS processing ends. Used to check for Ancestors.
    std::vector<int> subtree_size; // Number of vertices in subtree rooted at each vertex.
    std::vector<std::vector<int>> chains; // The i-th vector is the i-th chain
    std::vector<int> chain_index; // Which chain is a vertex part of? (Starts from 0)
    std::vector<int> position_in_chain; // How deep is a vertex in the chain? (0 is closest to the root)

    // DP Variables
    std::vector<std::vector<path>> lca_paths; // The i-th vector represents all paths that has i as the LCA
    std::vector<std::vector<long long>> cumulative_chain_children_memo; // The j-th index of the i-th vector represents the sum of the children's DP values for the j-th vertex and beyond within the i-th chain
    std::vector<std::vector<long long>> cumulative_chain_vertex_memo; // Exactly the same as above, but stores the vertex's DP value instead of the children's

    void dfs(int u, int& time) {
        dfs_start[u] = time;
        subtree_size[u] = 1;

        for (int v : adj_list[u]) {
            if (v == parent[u]) {
                continue;
            }

            parent[v] = u;
            depth[v] = depth[u] + 1;

            ++time;
            dfs(v, time);

            subtree_size[u] += subtree_size[v];
        }

        dfs_end[u] = time;
        ++time;
    }

    void make_rooted_tree() {
        parent[ROOT] = ROOT;
        depth[ROOT] = 0;

        int time = 0;
        dfs(ROOT, time);
    }

    void start_new_chain(int head) {
        int new_chain_id = chains.size();

        chains.push_back({head});
        chain_index[head] = new_chain_id;
        position_in_chain[head] = 0;
    }

    void recursive_construct_chain(int u) {
        int next_in_chain = UNCHECKED;
        for (int v : adj_list[u]) {
            if (v == parent[u]) {
                continue;
            }

            if (next_in_chain == UNCHECKED || subtree_size[next_in_chain] < subtree_size[v]) {
                next_in_chain = v;
            }
        }

        // Case: No children
        if (next_in_chain == UNCHECKED) {
            return;
        }

        // Set up next vertex in chain
        chains[chain_index[u]].push_back(next_in_chain);
        chain_index[next_in_chain] = chain_index[u];
        position_in_chain[next_in_chain] = position_in_chain[u] + 1;
        recursive_construct_chain(next_in_chain);

        // Start new chains in all other children
        for (int v : adj_list[u]) {
            if (v == parent[u] || v == next_in_chain) {
                continue;
            }

            start_new_chain(v);
            recursive_construct_chain(v);
        }
    }

    void construct_chains() {
        start_new_chain(ROOT);
        recursive_construct_chain(ROOT);
    }

    // Returns true if u is an ancestor of v
    bool is_ancestor(int u, int v) {
        return dfs_start[u] <= dfs_start[v] && dfs_end[v] <= dfs_end[u];
    }

    int find_lca(int u, int v) {
        // We don't really need HLD to find the LCA. 
        // This can be done in O(1) per query using a Sparse Table.
        // (But I didn't know at the time, so oh well)
        
        // Step 1: Find chain containing LCA
        int u_chain_index = chain_index[u];
        int u_curr = chains[u_chain_index][0];

        while (!is_ancestor(u_curr, v)) {
            u_curr = parent[u_curr]; // Move to parent chain
            u_chain_index = chain_index[u_curr];
            u_curr = chains[u_chain_index][0]; // Move to head of chain
        }

        int lca_chain_index = chain_index[u_curr];

        int low = 0;
        int high = (int) (chains[lca_chain_index].size()) - 1;
        int lca = u_curr;

        // Step 2: Perform Binary Search on the chain
        while (low <= high) {
            int mid = low + (high - low) / 2;
            int vertex = chains[lca_chain_index][mid];

            if (is_ancestor(vertex, v) && is_ancestor(vertex, u)) {
                lca = vertex;
                low = mid + 1;
            } else {
                high = mid - 1;
            }
        }

        return lca;
    }

    long long chain_sum(chain_value_type type, int chain_index, int chain_pos_start, int chain_pos_end) {
        auto &cumulative_memo = (type == chain_value_type::children 
                ? cumulative_chain_children_memo 
                : cumulative_chain_vertex_memo);
        
        auto &chain = cumulative_memo[chain_index];
        return chain[chain_pos_start] - chain[chain_pos_end + 1];
    }

    long long process_path(const path& p) {
        /*
        Here's the idea behind how this works. 
        
        Consider the following tree, and the path from A1 to A5 (i.e. A1-A2-A3-A4-A5):
        
                         A3
                       /    \
                     A2      A4
                    /  \    /  \
                   A1  X3  X4   A5
                  / |          / | \
                 X1 X2        X5 X6 X7
        
        Let DP[V] represent the highest possible score (i.e. number of votes) obtainable
        among all the paths that exist in the subtree rooted at the vertex V.
        
        If we decide to take the path from A1 to A5, then the maximum remaining score we can
        obtain is the sum of the DP values of the remaining untouched children. In this case, 
        that would be
        
        score(A1-A2-A3-A4-A5) + DP[X1] + DP[X2] + DP[X3] + DP[X4] + DP[X5] + DP[X6] + DP[X7]
        
        The DP terms can simply be added together because all the subtrees rooted at those 
        vertices are disjoint.
        
        This is still hard to deal with, since there may be many children to consider, so 
        let's try to represent it in a different way.
        
        Let children(V) represent the set of children of a vertex. For example, children(A2)
        returns {A1, X3}. Also for simplicity, let DP[{U1, ... , Uk}] = DP[U1]  + ... + DP[Uk].
        
        Now, we can represent the expression as such:
        
        score(A1-A2-A3-A4-A5) 
        + DP[children(A1)] + DP[children(A2)] + DP[children(A3)] + DP[children(A4)] + DP[children(A5)]
        - DP[A1] - DP[A2] - DP[A4] - DP[A5]
        
        Or, to make the expression more general, for a path P = A-...-L-...-B, where L is the 
        LCA of vertices A and B:
        
        score(P) + sum(DP[children(v)] for v in P) - sum(DP[v] for v in P excluding L)
        */

        int lca_chain_index = chain_index[p.lca];
        int lca_pos_in_chain = position_in_chain[p.lca];

        long long answer = 0;

        int endpoints[] = {p.u, p.v};
        for (int vertex : endpoints) {
            // Step 1: Climb up from curr to LCA
            long long path_value = 0;

            // Step 1a: Climb until the same chain as LCA
            int curr = vertex;
            while (chain_index[curr] != lca_chain_index) {
                int curr_chain_index = chain_index[curr];
                int curr_pos_in_chain = position_in_chain[curr];
                int curr_chain_head = chains[curr_chain_index][0];

                path_value += chain_sum(chain_value_type::children, curr_chain_index, 0, curr_pos_in_chain);
                // cumulative_chain_children_memo[curr_chain_index][0] - cumulative_chain_children_memo[curr_chain_index][curr_pos_in_chain + 1];
                path_value -= chain_sum(chain_value_type::vertex, curr_chain_index, 0, curr_pos_in_chain);
                // path_value -= cumulative_chain_vertex_memo[curr_chain_index][0] - cumulative_chain_vertex_memo[curr_chain_index][curr_pos_in_chain + 1];

                curr = parent[curr_chain_head];
            }

            // Step 1b: Climb from curr to LCA
            {
                int curr_pos_in_chain = position_in_chain[curr];
                path_value += chain_sum(chain_value_type::children, lca_chain_index, lca_pos_in_chain + 1, curr_pos_in_chain);
                path_value -= chain_sum(chain_value_type::vertex, lca_chain_index, lca_pos_in_chain + 1, curr_pos_in_chain);
                // path_value += cumulative_chain_children_memo[lca_chain_index][lca_pos_in_chain + 1] - cumulative_chain_children_memo[lca_chain_index][curr_pos_in_chain + 1];
                // path_value -= cumulative_chain_vertex_memo[lca_chain_index][lca_pos_in_chain + 1] - cumulative_chain_vertex_memo[lca_chain_index][curr_pos_in_chain + 1];
            }

            answer += path_value;
        }

        answer += chain_sum(chain_value_type::children, lca_chain_index, lca_pos_in_chain, lca_pos_in_chain) + p.value;

        // answer += (cumulative_chain_children_memo[lca_chain_index][lca_pos_in_chain] - cumulative_chain_children_memo[lca_chain_index][lca_pos_in_chain + 1]) + p.value;
        return answer;
    }

    long long bottom_up_dp(int u) {
        // DP value of u represents the maximum value obtainable within subtree rooted at u.

        // Step 1: Process all children first. Keep track of sum of DP values of children
        long long children_dp = 0;
        for (int v : adj_list[u]) {
            if (v == parent[u]) {
                continue;
            }

            children_dp += bottom_up_dp(v);
        }

        // Step 2: Update cumulative chain children memo
        int u_chain_index = chain_index[u];
        int u_pos_in_chain = position_in_chain[u];

        cumulative_chain_children_memo[u_chain_index][u_pos_in_chain] = cumulative_chain_children_memo[u_chain_index][u_pos_in_chain + 1] + children_dp;

        // Step 3: Process all paths to find DP value of u
        long long u_dp_value = children_dp;

        for (path& p : lca_paths[u]) {
            u_dp_value = std::max(u_dp_value, process_path(p));
        }

        // Step 4: Update cumulative chain vertex memo
        cumulative_chain_vertex_memo[u_chain_index][u_pos_in_chain] = cumulative_chain_vertex_memo[u_chain_index][u_pos_in_chain + 1] + u_dp_value;

        // Step 5: Return
        return u_dp_value;
    }

public:
    heavy_light(int _num_vertices) :
            num_vertices(_num_vertices),
            adj_list(num_vertices),
            parent(num_vertices, UNCHECKED),
            depth(num_vertices, UNCHECKED),
            dfs_start(num_vertices, UNCHECKED),
            dfs_end(num_vertices, UNCHECKED),
            subtree_size(num_vertices, UNCHECKED),
            chain_index(num_vertices, UNCHECKED),
            position_in_chain(num_vertices, UNCHECKED),
            lca_paths(num_vertices) {}

    void add_edge(int u, int v) {
        adj_list[u].push_back(v);
        adj_list[v].push_back(u);
    }

    void decompose() {
        make_rooted_tree();
        construct_chains();
    }

    void print_details() {
        for (int i = 0; i < (int) chains.size(); ++i) {
            std::cout << "Chain #" << i << ": ";
            for (int u : chains[i]) {
                std::cout << u << "-";
            }
            std::cout << "\n";
        }
        std::cout << "\n";

        for (int i = 0; i < num_vertices; ++i) {
            std::cout << "Vertex #" << i
                    << " - DFS times: " << dfs_start[i] << " to " << dfs_end[i]
                    << ", Chain ID: " << chain_index[i]
                    << ", Position in Chain: " << position_in_chain[i] << "\n";
        }
    }
    void add_path(int u, int v, long long value) {
        int lca = find_lca(u, v);
        lca_paths[lca].push_back({u, v, lca, value});
    }

    long long get_answer() {
        for (auto& chain : chains) {
            cumulative_chain_children_memo.emplace_back(chain.size() + 1, 0);
            cumulative_chain_vertex_memo.emplace_back(chain.size() + 1, 0);
        }

        return bottom_up_dp(ROOT);
    }
};

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int num_vertices;
    std::cin >> num_vertices;

    heavy_light hld(num_vertices);

    for (int i = 0; i < num_vertices - 1; ++i) {
        int u, v;
        std::cin >> u >> v;

        hld.add_edge(u - 1, v - 1); // 1-indexed vertices
    }

    hld.decompose();

    int num_paths;
    std::cin >> num_paths;

    for (int i = 0; i < num_paths; ++i) {
        int u, v;
        long long value;

        std::cin >> u >> v >> value;

        hld.add_path(u - 1, v - 1, value); // 1-indexed vertices
    }

    std::cout << hld.get_answer() << "\n";
}

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