Problem p02188 - Generation 2

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.
N-1 pairs of integers, each pair representing a connection between two entities, 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(inputs) that takes in the input.
```python
def f(inputs: List[Tuple[int, List[int], List[Tuple[int, int]]]]):
    ''' 
    inputs: a list of tuples, each containing: 
        - an integer 
        - a list of integers
        - a list of tuples, each containing two integers
    '''
```

Test Case Generator

Sure, here's a test case generator for the "G: Restricted DFS" problem statement:

```python
import random

def generate_tree(N):
    edges = []
    for i in range(2, N + 1):
        edges.append((random.randint(1, i - 1), i))
    return edges

def tcgen():
    # Generate a random number of vertices within the given constraints
    N = random.randint(1, 3 * 10**5)
    
    # Generate a list of N integers, each in the range [0, 10^9]
    A = [random.randint(0, 10**9) for _ in range(N)]
    
    # Generate N-1 edges to form a tree
    edges = generate_tree(N)
    
    return N, A, edges

# Example usage:
if __name__ == "__main__":
    N, A, edges = tcgen()
    print(N)
    print(" ".join(map(str, A)))
    for u, v in edges:
        print(u, v)
```

### Explanation:
1. **Number of Vertices (N)**: Randomly generates an integer \( N \) between 1 and 300,000.
2. **List of Integers (A)**: Creates a list of \( N \) random integers, each ranging from 0 to 1,000,000,000.
3. **Tree Edges**: To form a valid tree, we start from vertex 2 to \( N \) and connect each vertex to a random vertex from the previously created vertices. This approach ensures there are no cycles and the graph remains a tree.

### Example Output:
```
3
1 2 3
1 2
1 3
```

This ensures that you get a valid tree structure and adheres to the constraints given. You can call `tcgen()` to generate a variety of test cases for this problem.

Extract Arguments

def extract_arguments(fh):
    N = int(fh.readline().strip())
    A = list(map(int, fh.readline().strip().split()))
    edges = [tuple(map(int, fh.readline().strip().split())) for _ in range(N - 1)]
    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)