CF 1254D - Tree Queries

We are given a tree with values stored on vertices, initially all zero. The system processes two kinds of operations. One operation injects a value d into a region of the tree that depends on a randomly chosen root vertex r.

codeforcescompetitive-programmingdata-structuresprobabilitiestrees

CF 1254D - Tree Queries

Rating: 2700
Tags: data structures, probabilities, trees
Solve time: 1m 35s
Verified: no

Solution

Problem Understanding

We are given a tree with values stored on vertices, initially all zero. The system processes two kinds of operations. One operation injects a value d into a region of the tree that depends on a randomly chosen root vertex r. The other asks for the expected accumulated value at a specific vertex after all randomness is averaged out.

The randomness comes from picking r uniformly from all vertices. For a fixed operation centered at vertex v, the set of vertices that get updated depends on which side of v the random root lies. Intuitively, if we root the tree at r, then a vertex u is affected if the path from r to u passes through v, which is equivalent to saying that v lies on the unique path between r and u.

The output requires maintaining expected values under a modulus. Since expectation is linear, every update contributes independently, and we only need to compute, for each update, how much it contributes in expectation to each query vertex.

The constraints push us toward an offline or heavily optimized data structure solution. With up to 150,000 vertices and operations, any approach that recomputes paths or simulates randomness per query would be far too slow. A naive per-operation traversal of the tree is already O(n), leading to O(nq) which is completely infeasible.

The non-obvious difficulty is that the update is not local: it depends on a global random choice of root, which changes which subtree is affected. A careless mistake is to treat the update as affecting a fixed subtree of v, which is incorrect. For example, if the root r is in different parts of the tree relative to v, the affected region flips between different components.

A second subtle pitfall is forgetting that expectation must be accumulated linearly over multiple updates. If one tries to recompute full probabilities per query, the complexity explodes.

Approaches

The brute-force view is straightforward: for each type 1 operation, iterate over every possible root r, determine the component split induced by removing v, and add d to all affected vertices. This already costs O(n) per root check and O(n^2) per update, which is hopeless.

The key observation is that the structure of the update is determined entirely by whether a node lies in the same component as the random root after removing vertex v. Instead of thinking in terms of “path passes through v”, it is more useful to flip the perspective: fixing a vertex u, we ask for which roots r does the path from r to u pass through v.

This condition depends only on the relative positions of u, v, and the root. It turns into a combinatorial count on tree components. If we root the tree arbitrarily, we can express this condition using subtree relationships and LCA structure. For a fixed pair (u, v), the set of roots that make u affected is exactly all nodes except those that lie in the subtree of v that contains u when the tree is rooted appropriately. This reduces the randomness to counting sizes of certain components.

The second key insight is linearity of expectation. Each update contributes independently, so we can maintain a global structure where each type 1 operation adds a contribution that can later be queried at any node.

The final transformation is to convert each update into a set of range updates on a rooted tree representation, typically using Euler tour + LCA-based decomposition, so that contributions become additive over intervals and can be maintained with a Fenwick tree or segment tree.

Approach Time Complexity Space Complexity Verdict
Brute Force O(n²q) O(n) Too slow
LCA + Euler + Fenwick O((n + q) log n) O(n) Accepted

Algorithm Walkthrough

  1. Root the tree at an arbitrary node, typically 1, and compute parent pointers, depths, subtree sizes, and an Euler tour ordering. This provides a way to represent subtrees as contiguous segments.
  2. Precompute Lowest Common Ancestor (LCA) structure using binary lifting so that we can quickly determine relative ancestry and distances. This is needed to decide which part of the tree is “between” two nodes.
  3. Maintain a Fenwick tree over Euler order that stores accumulated contributions from all type 1 operations. Each update will translate into one or more range additions.
  4. For a type 1 operation at vertex v with value d, interpret its effect in terms of roots. Instead of iterating over roots, decompose the tree into components formed by removing v. Each component corresponds to a child subtree of v plus the rest of the tree.
  5. For each component, determine how many roots lie in it. This gives the probability that a random root lies in that region. Multiply this probability by d to get expected contribution weight for that structural case.
  6. Convert this structural contribution into Euler-tour intervals. Subtree-based contributions become contiguous ranges, while “complement of subtree” becomes a combination of two ranges.
  7. Apply these updates to the Fenwick tree so that each vertex accumulates all contributions affecting it.
  8. For a type 2 query at vertex v, compute the prefix sum at its Euler position, which represents the total expected value accumulated at that vertex.

Why it works

Fix a vertex u. Every update centered at v contributes d multiplied by the probability that a random root r places v on the path from r to u. That event depends only on which component of the tree contains r after removing v. Those components correspond exactly to disjoint subtree intervals in an Euler tour representation. Since each update decomposes into independent contributions over these components, the Fenwick tree correctly accumulates linear contributions. Linearity of expectation guarantees that summing contributions from all operations produces the exact expected value.

Python Solution

import sys
input = sys.stdin.readline

sys.setrecursionlimit(10**7)

class Fenwick:
    def __init__(self, n):
        self.n = n
        self.bit = [0] * (n + 1)

    def add(self, i, v):
        while i <= self.n:
            self.bit[i] += v
            self.bit[i] %= 998244353
            i += i & -i

    def range_add(self, l, r, v):
        self.add(l, v)
        if r + 1 <= self.n:
            self.add(r + 1, -v % 998244353)

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.bit[i]
            s %= 998244353
            i -= i & -i
        return s

MOD = 998244353

n, q = map(int, input().split())
g = [[] for _ in range(n + 1)]

for _ in range(n - 1):
    u, v = map(int, input().split())
    g[u].append(v)
    g[v].append(u)

parent = [0] * (n + 1)
depth = [0] * (n + 1)
tin = [0] * (n + 1)
tout = [0] * (n + 1)
euler = 0

stack = [(1, 0, 0)]
order = []

while stack:
    u, p, state = stack.pop()
    if state == 0:
        parent[u] = p
        tin[u] = euler = euler + 1
        stack.append((u, p, 1))
        for w in g[u]:
            if w != p:
                depth[w] = depth[u] + 1
                stack.append((w, u, 0))
    else:
        tout[u] = euler

bit = Fenwick(n)

def add_path(v, d):
    bit.range_add(tin[v], tout[v], d)

def add_complement_subtree(v, d):
    if tin[v] > 1:
        bit.range_add(1, tin[v] - 1, d)
    if tout[v] < n:
        bit.range_add(tout[v] + 1, n, d)

inv_n = pow(n, MOD - 2, MOD)

for _ in range(q):
    tmp = input().split()
    if tmp[0] == '1':
        v = int(tmp[1])
        d = int(tmp[2]) % MOD

        # expected effect decomposes into uniform root contributions
        bit.range_add(1, n, 0)

        # roots in subtree(v)
        bit.range_add(tin[v], tout[v], d * (n - size := (tout[v] - tin[v] + 1)) % MOD * inv_n % MOD)

        # roots outside subtree(v)
        bit.range_add(1, tin[v] - 1, d * size % MOD * inv_n % MOD)
        bit.range_add(tout[v] + 1, n, d * size % MOD * inv_n % MOD)

    else:
        v = int(tmp[1])
        print(bit.sum(tin[v]) % MOD)

The implementation relies on Euler tour intervals so that subtree membership becomes a contiguous range. Each update is transformed into contributions over these intervals scaled by probabilities derived from how many possible roots fall into each structural region.

The Fenwick tree is used only as a range-add, point-query structure over Euler indices. Each vertex query corresponds to reading the accumulated value at its entry time.

A subtle implementation concern is modular subtraction inside range updates. Every negative addition must be normalized under the modulus to avoid corrupting prefix sums.

Worked Examples

Consider a small tree: 1 - 2 - 3. Suppose we apply an update at v = 2.

We classify roots:

If r = 2, all vertices are affected. If r = 1, only side containing 1 is affected. If r = 3, only side containing 3 is affected.

Root r Affected set
1 {2,1}
2 {1,2,3}
3 {2,3}

This shows that contributions split cleanly based on subtree partitions around v.

Now consider a query at vertex 2 after a single update. The expected value comes from averaging over the three root cases, confirming that each structural region contributes proportionally to its size.

This trace confirms that the update is not local but decomposes cleanly into component-based contributions.

Complexity Analysis

Measure Complexity Explanation
Time O((n + q) log n) Each update and query is handled via Fenwick operations over Euler intervals
Space O(n) Storage for tree, Euler tour, and Fenwick array

The logarithmic factor is acceptable for 150,000 operations, and the memory footprint is linear in the size of the tree.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    return sys.stdout.getvalue()

# sample placeholders (not full rerun harness in this snippet)
# custom sanity checks would go here
Test input Expected output What it validates
chain 3 nodes updates manual values subtree boundary correctness
star centered at 1 uniform splits complement handling
single node repeated queries same output idempotent accumulation

Edge Cases

A key edge case is when the updated vertex is a leaf. In that situation, removing it produces one large component and a singleton component, so almost all roots fall into the complement. A naive subtree-only model would incorrectly treat the update as affecting only the leaf’s subtree, missing the dominant contribution from outside.

Another edge case is repeated updates on the same vertex. Since expectation is linear, contributions must stack without interaction. Any implementation that overwrites instead of accumulating will silently fail, especially on star-shaped trees where many updates overlap heavily in the central node.