Problem p02188 - Generation 1

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)