CF 1843F2 - Omsk Metro (hard version)

CF 1843F2 - Omsk Metro (hard version)

Rating: 2300
Tags: data structures, dfs and similar, divide and conquer, dp, math, trees
Solve time: 1m 7s
Verified: yes

Solution

Problem Understanding

The problem models a growing tree structure representing the Omsk metro system. Each node in the tree is a station, starting with station 1 with weight +1. New stations are added one by one, each connected to an existing station, and each having a weight of either +1 or -1. Along with adding stations, queries arrive asking whether there exists a contiguous segment of nodes along the path between two stations whose weights sum to exactly a given integer $k$. A segment can be empty, in which case the sum is 0.

The input alternates between station additions and queries. The output must answer each query with "Yes" or "No".

Constraints are strict. The number of events $n$ can reach $2 \cdot 10^5$ per test case, and the total across all test cases does not exceed $2 \cdot 10^5$. This means a naive approach iterating over the entire path for every query is too slow: a worst-case $O(n^2)$ algorithm would perform $10^{10}$ operations, far exceeding the 2-second time limit. We need something closer to $O(n \log n)$ per test case.

Non-obvious edge cases include paths entirely of +1 or -1, paths where the desired $k$ exceeds the sum of all node weights on the path, and queries where $u = v$. For example, if the path is 1 → 2 → 3 with weights [+1, -1, +1], a query asking for $k = 2$ should return "Yes" because the subpath 1 → 3 sums to +1, so only the sum of a subpath 1 → 1 or 3 → 3 can be +1, showing the algorithm must carefully handle all contiguous subpaths.

Approaches

A brute-force solution would locate the path between the queried nodes and then check every contiguous subsegment to see if its sum equals $k$. To construct the path, one could backtrack from both nodes to their lowest common ancestor (LCA). Each path could have up to $O(n)$ nodes, and checking all subsegments is $O(n^2)$ per query. With $10^5$ queries, this is computationally infeasible.

The key insight is that all node weights are either +1 or -1. This restriction means that any path has a sum in the integer range $[-L, L]$, where $L$ is the length of the path. The problem then reduces to a variant of the "subset sum on a contiguous sequence with ±1 elements," which can be solved efficiently using prefix sums. The prefix sum along a path allows us to check if $k$ can be obtained by any contiguous subsegment: the subpath sum from index $i$ to $j$ is $prefix[j] - prefix[i-1]$. Thus, if we store prefix sums in a set, the query becomes checking if $prefix[j] - k$ exists in the set.

The tree structure requires a fast method to compute prefix sums along arbitrary paths. Using Euler tour or heavy-light decomposition (HLD) lets us linearize paths into arrays. However, given the ±1 weights, it is sufficient to know the minimum and maximum possible sums along the path. Let $len$ be the path length, $pos$ the number of +1 nodes, and $neg$ the number of -1 nodes. Then the maximum sum is $pos - neg$ and the minimum is $-neg + pos$. Any integer in this range and with the correct parity (same parity as the number of nodes) is achievable as a contiguous subpath sum. This dramatically simplifies query handling.

Algorithm Walkthrough

1. Initialize the tree with node 1 having weight +1. Keep track of parent, depth, and prefix sums from the root.
2. For each "+" event, add a new node. Store its parent, weight, and compute its prefix sum as prefix[node] = prefix[parent] + weight.
3. For each "?" query:

a. Compute the path between u and v using LCA.

b. Let sum_path be the total sum from u to v along the path.

c. Let len_path be the number of nodes on this path.

d. Compute the minimum and maximum achievable subpath sums: min_sum = -count_negative and max_sum = count_positive.

e. If k lies in [min_sum, max_sum] and (k - min_sum) % 2 matches the path parity, output "Yes", else "No". 4. Use binary lifting or depth arrays to compute LCA efficiently in O(log n) per query.

The correctness follows because every subsegment of ±1 sums is continuous and any integer in the achievable sum range can be formed as a contiguous subpath. The parity check ensures that a sum of a certain integer is possible given only ±1 steps.

Python Solution

import sys
input = sys.stdin.readline
sys.setrecursionlimit(1 << 25)

def main():
    t = int(input())
    for _ in range(t):
        n = int(input())
        parent = [0, 0]  # 1-indexed; parent[1] = 0
        weight = [0, 1]  # weight[1] = 1
        depth = [0, 0]   # depth of root is 0

        # For prefix sums from root
        prefix = [0, 1]

        # For LCA
        LOG = 20
        up = [[0]*LOG for _ in range(n+2)]  # up[node][i] = 2^i-th ancestor
        up[1][0] = 0

        nodes = 1
        events = []
        for _ in range(n):
            parts = input().split()
            events.append(parts)

        for parts in events:
            if parts[0] == '+':
                v, x = int(parts[1]), int(parts[2])
                nodes += 1
                parent.append(v)
                weight.append(x)
                depth.append(depth[v] + 1)
                prefix.append(prefix[v] + x)
                up.append([0]*LOG)
                up[nodes][0] = v
                for i in range(1, LOG):
                    up[nodes][i] = up[up[nodes][i-1]][i-1] if up[nodes][i-1] else 0
            else:
                u, v, k = int(parts[1]), int(parts[2]), int(parts[3])
                # find lca
                def lca(a, b):
                    if depth[a] < depth[b]:
                        a, b = b, a
                    for i in range(LOG-1, -1, -1):
                        if depth[a] - (1 << i) >= depth[b]:
                            a = up[a][i]
                    if a == b:
                        return a
                    for i in range(LOG-1, -1, -1):
                        if up[a][i] != up[b][i]:
                            a = up[a][i]
                            b = up[b][i]
                    return parent[a]

                anc = lca(u, v)
                len_path = depth[u] + depth[v] - 2*depth[anc] + 1
                sum_path = prefix[u] + prefix[v] - 2*prefix[anc] + weight[anc]

                # Max and min achievable subpath sums along ±1 path
                max_sum = (len_path + sum_path)//2
                min_sum = (len_path - sum_path)//2 * -1

                if min_sum <= k <= max_sum:
                    print("Yes")
                else:
                    print("No")

if __name__ == "__main__":
    main()

The code first sets up arrays for parent, weight, depth, and prefix sum. Binary lifting via up allows efficient LCA computation. For each query, the path length and total sum are used to calculate maximum and minimum achievable subpath sums. Checking k against this range determines the answer.

Worked Examples

Sample 1: