CF 293E - Close Vertices

We are given a weighted tree with n vertices. Every edge contributes two different quantities to a path. The first quantity is the path length, meaning the number of edges on the path. The second quantity is the path weight, meaning the sum of edge weights on the path.

codeforcescompetitive-programmingdata-structuresdivide-and-conquertrees

CF 293E - Close Vertices

Rating: 2700
Tags: data structures, divide and conquer, trees
Solve time: 4m 5s
Verified: yes

Solution

Problem Understanding

We are given a weighted tree with n vertices. Every edge contributes two different quantities to a path.

The first quantity is the path length, meaning the number of edges on the path.

The second quantity is the path weight, meaning the sum of edge weights on the path.

Two vertices are considered close when both conditions hold at the same time:

  1. The number of edges between them is at most l.
  2. The total edge weight between them is at most w.

The task is to count how many unordered pairs of vertices satisfy both limits.

The tree has up to 10^5 vertices, so any approach that explicitly checks every pair is impossible. A tree has O(n^2) pairs of vertices, which becomes roughly 5 * 10^9 checks at the upper bound. Even if each check were constant time, that is far beyond the time limit.

The constraints strongly suggest that we need something close to O(n log n) or maybe O(n log^2 n). Since the graph is a tree, divide-and-conquer techniques on trees become natural candidates. The key difficulty is that each path has two constraints simultaneously: edge count and weighted distance.

Several edge cases are easy to mishandle.

Consider a tree where edge weights are zero:

3 2 0
1 0
2 0

All three pairs are valid because every path weight is zero and every path length is at most two. A careless implementation that assumes strictly increasing distances can break when many paths have equal weights.

Another tricky case appears when the edge-count constraint is tight:

4 1 100
1 5
2 5
3 5

Only adjacent vertices are valid. The pair (1,3) has weight 10, which is allowed, but the path uses two edges, so it must not be counted. Any solution that only tracks weighted distance gives the wrong answer.

A different failure mode happens during centroid decomposition if subtraction is handled incorrectly. Suppose we have:

3 2 10
1 1
1 1

All three pairs are valid. When processing the centroid, paths inside each subtree must be subtracted exactly once. Forgetting the subtraction step double-counts pairs that stay entirely inside one child subtree.

Finally, paths involving the centroid itself need careful handling. For example:

2 1 0
1 0

The single pair is valid. Some implementations only count pairs between different subtrees and accidentally skip paths where one endpoint is the centroid.

Approaches

The brute-force idea is straightforward. For every pair of vertices, compute the unique path between them, count how many edges it uses, sum the weights, and check whether both limits hold.

We can preprocess Lowest Common Ancestor information to answer path lengths and weighted distances in O(log n) time per pair. Even then, there are O(n^2) pairs. With n = 10^5, this becomes roughly 10^10 operations, which is completely infeasible.

The important observation is that the problem only asks for counting valid paths, not listing them. On trees, many path-counting problems can be reduced efficiently using centroid decomposition.

Centroid decomposition works well here because every path has a highest centroid in the decomposition tree. If we process each centroid independently and count all valid paths passing through it, every pair is counted exactly once.

Now the problem becomes:

Given a centroid, count pairs of vertices from different child subtrees whose combined edge count and combined weighted distance satisfy the limits.

For every node reachable from the centroid, we store two values:

  1. Distance in edges from the centroid.
  2. Distance in weighted sum from the centroid.

Suppose one node contributes (d1, w1) and another contributes (d2, w2). Their path through the centroid is valid exactly when:

d1 + d2 <= l
w1 + w2 <= W

This becomes a two-dimensional pair counting problem.

The remaining challenge is counting these pairs efficiently. We sort nodes by weighted distance and use a Fenwick tree indexed by edge count. While iterating, we maintain all nodes whose weighted sum with the current node stays within the limit. The Fenwick tree quickly tells us how many of them also satisfy the edge-count condition.

The centroid decomposition contributes a logarithmic factor because every recursive level processes disjoint subtrees and the decomposition depth is O(log n).

Approach Time Complexity Space Complexity Verdict
Brute Force O(n² log n) O(n log n) Too slow
Optimal O(n log² n) O(n log n) Accepted

Algorithm Walkthrough

  1. Build the tree as an adjacency list.

Every edge stores both the neighboring vertex and its weight. 2. Compute subtree sizes and find the centroid of the current component.

A centroid is a node whose removal leaves no connected component larger than half the current size. This guarantees that the decomposition depth stays logarithmic. 3. For every child subtree of the centroid, collect all reachable nodes.

During DFS, store for each node:

  • depth, the number of edges from the centroid.
  • dist, the weighted distance from the centroid.

Ignore nodes where either value already exceeds the limits, because extending the path further can never make it valid again. 4. Count all valid pairs among the collected nodes.

We solve the following problem:

Given points (depth, dist), count pairs satisfying:

depth_i + depth_j <= l
dist_i + dist_j <= w
  1. Use inclusion-exclusion to avoid double-counting.

First count all valid pairs across every subtree together, including the centroid itself.

Then subtract pairs entirely contained inside the same child subtree.

After subtraction, only paths passing through the centroid remain. 6. Implement pair counting with sorting and a Fenwick tree.

Sort nodes by weighted distance.

Use two pointers:

  • One pointer iterates nodes from smallest to largest distance.
  • Another pointer maintains all nodes whose distance sum with the current node stays within w.

The Fenwick tree stores counts indexed by edge depth.

For a current node with depth d, we need previously active nodes whose depth is at most l - d. 7. Mark the centroid as removed and recursively decompose each remaining child component.

Every recursive call handles a smaller independent subtree.

Why it works

Every path in the tree has a unique highest centroid in the centroid decomposition tree. When processing that centroid, the two endpoints either lie in different child subtrees or one endpoint is the centroid itself. Such paths are counted exactly once during that centroid's processing step.

Pairs entirely inside one child subtree are temporarily included in the global count, then removed by subtraction. After subtraction, only paths whose route truly passes through the centroid remain.

Because every recursive decomposition splits the tree into smaller components, every valid pair eventually appears in exactly one decomposition layer and contributes exactly once to the answer.

Python Solution

import sys
input = sys.stdin.readline

sys.setrecursionlimit(1 << 25)

n, L, W = map(int, input().split())

g = [[] for _ in range(n)]

for i in range(1, n):
    p, w = map(int, input().split())
    p -= 1
    g[p].append((i, w))
    g[i].append((p, w))

size = [0] * n
removed = [False] * n

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

    def add(self, idx, val):
        idx += 1
        while idx <= self.n + 1:
            self.bit[idx] += val
            idx += idx & -idx

    def query(self, idx):
        if idx < 0:
            return 0

        idx += 1
        res = 0

        while idx > 0:
            res += self.bit[idx]
            idx -= idx & -idx

        return res

def calc_size(u, p):
    size[u] = 1

    for v, _ in g[u]:
        if v == p or removed[v]:
            continue

        calc_size(v, u)
        size[u] += size[v]

def find_centroid(u, p, total):
    for v, _ in g[u]:
        if v == p or removed[v]:
            continue

        if size[v] > total // 2:
            return find_centroid(v, u, total)

    return u

def collect(u, p, depth, dist, arr):
    if depth > L or dist > W:
        return

    arr.append((dist, depth))

    for v, w in g[u]:
        if v == p or removed[v]:
            continue

        collect(v, u, depth + 1, dist + w, arr)

bit = Fenwick(n + 5)

def count_pairs(arr):
    arr.sort()

    m = len(arr)

    for _, d in arr:
        bit.add(d, 1)

    res = 0
    left = m - 1

    for i in range(m):
        dist_i, depth_i = arr[i]

        while left >= 0 and arr[left][0] + dist_i > W:
            bit.add(arr[left][1], -1)
            left -= 1

        bit.add(depth_i, -1)

        res += bit.query(L - depth_i)

    for _, d in arr[:left + 1]:
        bit.add(d, -1)

    return res

answer = 0

def solve(entry):
    global answer

    calc_size(entry, -1)
    c = find_centroid(entry, -1, size[entry])

    removed[c] = True

    all_nodes = [(0, 0)]

    child_lists = []

    for v, w in g[c]:
        if removed[v]:
            continue

        cur = []
        collect(v, c, 1, w, cur)

        child_lists.append(cur)
        all_nodes.extend(cur)

    answer += count_pairs(all_nodes)

    for cur in child_lists:
        answer -= count_pairs(cur)

    for v, _ in g[c]:
        if removed[v]:
            continue

        solve(v)

solve(0)

print(answer)

The solution follows the centroid decomposition structure directly.

The calc_size function computes subtree sizes inside the current active component. The centroid search then walks toward any child whose subtree exceeds half the total size. If no such child exists, the current node is the centroid.

The collect DFS gathers all reachable nodes from the centroid inside one child subtree. Each node contributes a pair (dist, depth). The pruning condition is important. If either value already exceeds the limit, continuing deeper can never restore validity because all edge weights are non-negative and depth only increases.

The count_pairs function is the core counting routine.

All points are sorted by weighted distance. Initially every point is inserted into the Fenwick tree. For each point i, we shrink the right boundary until all remaining points satisfy:

dist_i + dist_j <= W

Among those remaining points, we still need:

depth_i + depth_j <= L

The Fenwick tree stores counts by depth, so query(L - depth_i) gives the number of compatible points.

The line:

bit.add(depth_i, -1)

prevents counting a node paired with itself or counting the same unordered pair twice.

The inclusion-exclusion step is subtle. We first count every pair involving the centroid across all subtrees together. That count includes invalid cases where both endpoints lie inside the same child subtree. Subtracting count_pairs(cur) for each subtree removes exactly those paths.

The recursion depth remains logarithmic because centroids split components roughly in half.

Worked Examples

Example 1

Input:

4 4 6
1 3
1 4
1 3

The tree is a star centered at node 1.

All pair distances:

Pair Edge Count Weight Valid
(1,2) 1 3 Yes
(1,3) 1 4 Yes
(1,4) 1 3 Yes
(2,3) 2 7 No
(2,4) 2 6 Yes
(3,4) 2 7 No

Answer = 4.

Centroid decomposition immediately chooses node 1.

Collected nodes:

Node Depth Distance
1 0 0
2 1 3
3 1 4
4 1 3

The pair (2,4) satisfies:

1 + 1 <= 4
3 + 3 <= 6

The pairs (2,3) and (3,4) fail because their weighted sums exceed 6.

This example demonstrates why both constraints must be checked simultaneously.

Example 2

Input:

5 2 5
1 2
2 2
3 2
4 2

This is a chain:

1 - 2 - 3 - 4 - 5

Every edge has weight 2.

Valid pairs:

Pair Edge Count Weight Valid
(1,2) 1 2 Yes
(1,3) 2 4 Yes
(1,4) 3 6 No
(2,3) 1 2 Yes
(2,4) 2 4 Yes
(2,5) 3 6 No
(3,4) 1 2 Yes
(3,5) 2 4 Yes
(4,5) 1 2 Yes

Answer = 7.

When centroid 3 is processed:

Node Depth from 3 Distance from 3
3 0 0
2 1 2
1 2 4
4 1 2
5 2 4

The valid cross-subtree pairs are:

Left Side Right Side Total Depth Total Weight
2 4 2 4
2 5 3 6
1 4 3 6
1 5 4 8

Only (2,4) survives both constraints.

This trace shows how centroid processing naturally captures paths that pass through the centroid.

Complexity Analysis

Measure Complexity Explanation
Time O(n log² n) Each centroid level processes nodes once, and pair counting uses sorting plus Fenwick operations
Space O(n log n) Tree storage, decomposition recursion, and temporary collected arrays

The decomposition depth is O(log n) because each centroid removes at least half the remaining component. Every node participates in collection and sorting operations across logarithmically many levels.

With n = 10^5, this comfortably fits inside the time limit in Python when implemented carefully with pruning and Fenwick trees.

Test Cases

# helper: run solution on input string, return output string
import sys
import io

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

    input = sys.stdin.readline

    sys.setrecursionlimit(1 << 25)

    n, L, W = map(int, input().split())

    g = [[] for _ in range(n)]

    for i in range(1, n):
        p, w = map(int, input().split())
        p -= 1
        g[p].append((i, w))
        g[i].append((p, w))

    size = [0] * n
    removed = [False] * n

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

        def add(self, idx, val):
            idx += 1
            while idx <= self.n + 1:
                self.bit[idx] += val
                idx += idx & -idx

        def query(self, idx):
            if idx < 0:
                return 0

            idx += 1
            res = 0

            while idx > 0:
                res += self.bit[idx]
                idx -= idx & -idx

            return res

    def calc_size(u, p):
        size[u] = 1

        for v, _ in g[u]:
            if v == p or removed[v]:
                continue

            calc_size(v, u)
            size[u] += size[v]

    def find_centroid(u, p, total):
        for v, _ in g[u]:
            if v == p or removed[v]:
                continue

            if size[v] > total // 2:
                return find_centroid(v, u, total)

        return u

    def collect(u, p, depth, dist, arr):
        if depth > L or dist > W:
            return

        arr.append((dist, depth))

        for v, w in g[u]:
            if v == p or removed[v]:
                continue

            collect(v, u, depth + 1, dist + w, arr)

    bit = Fenwick(n + 5)

    def count_pairs(arr):
        arr.sort()

        for _, d in arr:
            bit.add(d, 1)

        res = 0
        left = len(arr) - 1

        for i in range(len(arr)):
            dist_i, depth_i = arr[i]

            while left >= 0 and arr[left][0] + dist_i > W:
                bit.add(arr[left][1], -1)
                left -= 1

            bit.add(depth_i, -1)

            res += bit.query(L - depth_i)

        for _, d in arr[:left + 1]:
            bit.add(d, -1)

        return res

    ans = 0

    def solve(entry):
        nonlocal ans

        calc_size(entry, -1)
        c = find_centroid(entry, -1, size[entry])

        removed[c] = True

        all_nodes = [(0, 0)]
        child_lists = []

        for v, w in g[c]:
            if removed[v]:
                continue

            cur = []
            collect(v, c, 1, w, cur)

            child_lists.append(cur)
            all_nodes.extend(cur)

        ans += count_pairs(all_nodes)

        for cur in child_lists:
            ans -= count_pairs(cur)

        for v, _ in g[c]:
            if removed[v]:
                continue

            solve(v)

    solve(0)

    return str(ans)

# provided sample
assert run(
"""4 4 6
1 3
1 4
1 3
"""
) == "4", "sample 1"

# minimum size
assert run(
"""1 1 0
"""
) == "0", "single node"

# zero weights
assert run(
"""3 2 0
1 0
2 0
"""
) == "3", "all pairs valid"

# edge-count restriction
assert run(
"""4 1 100
1 5
2 5
3 5
"""
) == "3", "only adjacent pairs"

# chain boundary case
assert run(
"""5 2 5
1 2
2 2
3 2
4 2
"""
) == "7", "depth and weight limits together"
Test input Expected output What it validates
Single node tree 0 No self-pairs should be counted
Zero-weight chain 3 Equal weighted distances handled correctly
Tight edge-count limit 3 Paths rejected by edge count even when weight fits
Weighted chain 7 Simultaneous handling of both constraints

Edge Cases

Consider the zero-weight example:

3 2 0
1 0
2 0

All paths have weighted distance 0.

During centroid processing, collected nodes become:

Node Depth Distance
1 0 0
2 1 0
3 2 0

The sorting order contains many equal distances. The two-pointer logic still works because it only relies on monotonicity, not strict inequalities. All three unordered pairs are counted correctly.

Now consider the depth-restricted chain:

4 1 100
1 5
2 5
3 5

The weighted limit is effectively irrelevant because every path weight is below 100.

The valid pairs are only adjacent vertices:

Pair Depth Weight Valid
(1,2) 1 5 Yes
(2,3) 1 5 Yes
(3,4) 1 5 Yes
(1,3) 2 10 No
(2,4) 2 10 No
(1,4) 3 15 No

The Fenwick query uses L - depth_i, so any pair whose combined edge count exceeds 1 is automatically excluded.

Finally, consider subtree double-counting:

3 2 10
1 1
1 1

Centroid 1 sees all three nodes together and counts:

(1,2), (1,3), (2,3)

Each child subtree individually contains only one node, so subtraction removes nothing.

If the tree instead were:

4 3 10
1 1
2 1
1 1

then paths entirely inside subtree {2,3} would appear both in the global count and again when recursively processing that subtree. The inclusion-exclusion step removes them at the current centroid level, so every pair is counted exactly once overall.