Orig Description
G: Restricted DFS
Problem
Given an undirected tree G with N vertices and N-1 edges, without self-loops or multiple edges. Each vertex is labeled with a non-negative integer A_i, and the vertices are numbered from 1 to N. The edges are also numbered from 1 to N-1, and the ith edge connects vertices u_i and v_i.
Consider performing depth-first search (DFS) on this tree, starting from the root r and following the pseudocode below.
// [input]
// G: The graph to be searched by DFS
// A: Non-negative integers assigned to each vertex
// v: The vertex to start DFS from
// step: An integer that records the number of steps taken
// [output]
// One of the following two values:
// - SUCCESS: DFS returns to vertex v without stopping in the middle
// - FAILURE: DFS stops in the middle
function dfs(G, A, v, step)
if (A[v] is 0) then
return FAILURE
A[v] ← A[v] - 1
step ← step + 1
Sort the children of v in ascending order of their vertex numbers
// c is visited in ascending order of vertex numbers
for each (child c of v) do
if (dfs(G, A, c, step) is FAILURE) then
return FAILURE
if (A[v] is 0) then
return FAILURE
A[v] ← A[v] - 1
step ← step + 1
return SUCCESS
That is, given G and A, execute
dfs(G, A, r, 0)
We want to find the number of steps taken by this DFS when each vertex is taken as the root.
Input
N
A_1 ... A_N
u_1 v_1
...
u_{N-1} v_{N-1}
The first line contains the number of vertices N of the graph.
The second line contains N integers. The ith integer A_i represents the value on the ith vertex.
The next N-1 lines contain information about the edges of the graph. Each line contains two integers u_i and v_i, indicating that there is an undirected edge between vertex u_i and vertex v_i.
Constraints
1 \leq N \leq 3 \times 10^5
0 \leq A_i \leq 10^9
1 \leq u_i < v_i \leq N
The graph given is guaranteed to be a tree.
All input is given as integers.
Output
Output N lines. The ith line should contain the number of steps taken in the DFS when vertex i is used as the root.
Sample Input 1
3
1 2 3
1 2
1 3
Sample Output 1
2
3
3
When vertex 1 is used as the root
Extracted Specification
An integer N (1 ≤ N ≤ 3 * 10^5), representing some quantity or size.
A list of N integers, where each integer is between 0 and 10^9.
A list of N-1 pairs of integers, where each integer in the pair is between 1 and N.
### Example Input:
```
3
1 2 3
1 2
1 3
```
### Function Signature:
Write a function f(N, A, edges) that takes in the input.
def f(N: int, A: List[int], edges: List[Tuple[int, int]]):
'''
N: an integer
A: a list of integers
edges: a list of pairs of integers
'''
Test Case Generator
Certainly! Here is a test case generator for the "Restricted DFS" problem:
```python
import random
def tcgen():
# Generate number of vertices N
N = random.randint(1, 3 * 10**5)
# Generate array A of size N with values in range [0, 10^9]
A = [random.randint(0, 10**9) for _ in range(N)]
# Generate N-1 edges to form a tree
edges = []
available_vertices = list(range(2, N + 1))
random.shuffle(available_vertices)
# Connecting node 1 to other nodes to ensure it's a tree
for i in range(2, N + 1):
parent = random.randint(1, i - 1)
edges.append((parent, i))
# Shuffle edges to ensure randomness in their order
random.shuffle(edges)
# Convert edges to the required input format
edges_str = "\n".join(f"{u} {v}" for u, v in edges)
# Format the entire input as a single string
input_str = f"{N}\n{' '.join(map(str, A))}\n{edges_str}"
return input_str
# Example of using the generator
print(tcgen())
```
### Explanation:
1. **Number of Vertices (N):**
- Randomly choose \( N \) from the range [1, 300000] to meet the problem constraints.
2. **Array A:**
- Generate an array \( A \) of length \( N \), where each value \( A[i] \) is a random integer between 0 and \( 10^9 \).
3. **Tree Edges:**
- Since the graph is a tree, it should have exactly \( N-1 \) edges.
- Ensure the tree structure by connecting vertex 1 (the root) to other vertices.
- Shuffle the vertex list to create randomness.
4. **Output Formatting:**
- Combine the number of vertices, the array \( A \), and the edges into a single string in the required format.
### Usage:
You can use the `tcgen` function to generate test cases for the "Restricted DFS" problem and validate your solution against those test cases.
Extract Arguments
def extract_arguments(fh):
N = int(fh.readline().strip())
A = list(map(int, fh.readline().strip().split()))
edges = []
for _ in range(N-1):
u, v = map(int, fh.readline().strip().split())
edges.append((u, v))
return N, A, edges
# if __name__ == "__main__":
# input_path = sys.argv[1]
# with open(input_path, 'r') as fh:
# N, A, edges = extract_arguments(fh)
# f(N, A, edges)