QTree Sum between Nodes

Competitive Programming
1

The “QTREE” problem on SPOJ involves answering queries about paths in a tree. Specifically, the problem typically requires calculating the sum of values associated with the nodes along the path between two nodes in a tree.

Problem Overview

Given:

  • A tree with N nodes.
  • Each node has a value.
  • You need to efficiently answer multiple queries asking for the sum of values along the path between two nodes u and v.

Key Concepts

  1. Tree Representation: Use an adjacency list to represent the tree.

  2. Lowest Common Ancestor (LCA): The LCA of two nodes is the deepest node that is an ancestor to both nodes. Finding the LCA helps break the path into segments that can be summed easily:
    [
    \text{Sum}(u, v) = \text{Sum}(1, u) + \text{Sum}(1, v) - 2 \times \text{Sum}(1, \text{LCA}(u, v)) + \text{Value}(\text{LCA}(u, v))
    ]

  3. DFS for Preprocessing: Perform a DFS traversal to preprocess information needed for answering queries, such as parent nodes and depths.

  4. Binary Lifting: This technique allows for efficient LCA queries in logarithmic time after a linear time preprocessing phase.

C++ Implementation

Here’s a step-by-step implementation of the QTREE problem:

#include <iostream>
#include <vector>
#include <cmath>

using namespace std;

const int MAXN = 100005; // Maximum number of nodes
const int LOG = 20;      // Sufficient for N <= 10^5

vector<int> tree[MAXN];
int parent[MAXN][LOG];
int depth[MAXN];
long long sumToRoot[MAXN];
long long nodeValue[MAXN]; // Store values associated with nodes

void dfs(int node, int par, long long currentSum) {
    parent[node][0] = par; // Set the immediate parent
    depth[node] = depth[par] + 1; // Set the depth of the current node
    sumToRoot[node] = currentSum + nodeValue[node]; // Calculate sum to root

    for (int child : tree[node]) {
        if (child != par) {
            dfs(child, node, sumToRoot[node]); // Recur for child nodes
        }
    }
}

// Preprocessing for Binary Lifting
void preprocess(int N) {
    for (int j = 1; j < LOG; ++j) {
        for (int i = 1; i <= N; ++i) {
            if (parent[i][j - 1] != -1) {
                parent[i][j] = parent[parent[i][j - 1]][j - 1];
            }
        }
    }
}

// Function to find the LCA of u and v
int lca(int u, int v) {
    if (depth[u] < depth[v]) {
        swap(u, v);
    }

    // Lift u to the same depth as v
    for (int j = LOG - 1; j >= 0; --j) {
        if (depth[parent[u][j]] >= depth[v]) {
            u = parent[u][j];
        }
    }

    // If they meet, return the LCA
    if (u == v) return u;

    for (int j = LOG - 1; j >= 0; --j) {
        if (parent[u][j] != parent[v][j]) {
            u = parent[u][j];
            v = parent[v][j];
        }
    }

    return parent[u][0]; // Return the LCA
}

int main() {
    int N, Q;
    cin >> N;

    // Read node values
    for (int i = 1; i <= N; ++i) {
        cin >> nodeValue[i];
    }

    // Read tree edges
    for (int i = 0; i < N - 1; ++i) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }

    // Initialize
    for (int i = 1; i <= N; ++i) {
        fill(parent[i], parent[i] + LOG, -1);
    }

    // Perform DFS to populate parent, depth, and sumToRoot
    dfs(1, 0, 0);
    preprocess(N);

    // Read queries
    cin >> Q;
    while (Q--) {
        int u, v;
        cin >> u >> v;
        int ancestor = lca(u, v);
        long long result = sumToRoot[u] + sumToRoot[v] - 2 * sumToRoot[ancestor] + nodeValue[ancestor];
        cout << result << endl;
    }

    return 0;
}

Explanation of the Code

  1. Data Structures:

    • tree: An adjacency list to store the tree.
    • parent: A 2D array for binary lifting to store ancestors.
    • depth: An array to store the depth of each node.
    • sumToRoot: An array to keep track of the cumulative sum of node values from the root to each node.
    • nodeValue: Stores values for each node.

  2. DFS Function:

    • Initializes parent, depth, and sumToRoot during the tree traversal.
    • Recursively processes child nodes.

  3. Binary Lifting Preprocessing:

    • Prepares the parent array for efficient LCA queries.

  4. LCA Function:

    • Computes the LCA using the binary lifting technique.

  5. Main Function:

    • Reads input values and tree structure.
    • Processes queries and computes the required sums using the previously calculated data.

Complexity

  • Preprocessing: (O(N \log N))
  • Querying: (O(\log N)) per query.
  • Total time for Q queries: (O(Q \log N)).

This solution efficiently handles the tree and query structure, allowing for quick calculations. Let me know if you have any questions or need further explanations!