/ SeriousOJ /

Record Detail

Accepted


  

Code

from collections import deque, defaultdict
import sys
input = sys.stdin.read

def bfs(start_node, graph, N):
    distances = [-1] * N
    distances[start_node] = 0
    queue = deque([start_node])
    
    max_dist = 0
    furthest_node = start_node
    
    while queue:
        node = queue.popleft()
        current_dist = distances[node]
        
        for neighbor in graph[node]:
            if distances[neighbor] == -1:
                distances[neighbor] = current_dist + 1
                queue.append(neighbor)
                if distances[neighbor] > max_dist:
                    max_dist = distances[neighbor]
                    furthest_node = neighbor
                    
    return furthest_node, max_dist

def minimum_seconds_to_collect_apples(T, test_cases):
    results = []
    
    for N, apple_nodes, edges in test_cases:
        if not any(apple_nodes):
            results.append("0")
            continue
        
        # Build the adjacency list for the tree
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        # Find any node with an apple
        start_node = next(i for i in range(N) if apple_nodes[i])
        
        # First BFS to find the farthest node from any apple node
        furthest_node, _ = bfs(start_node, graph, N)
        
        # Second BFS to find the maximum distance from this furthest node
        _, max_distance = bfs(furthest_node, graph, N)
        
        # The minimum time required is the maximum distance found
        results.append(str(max_distance))
    
    print("\n".join(results))

# Sample input processing
data = input().split()
T = int(data[0])
index = 1

test_cases = []
for _ in range(T):
    N = int(data[index])
    index += 1
    apple_nodes = list(map(int, data[index:index + N]))
    index += N
    edges = []
    for _ in range(N - 1):
        u = int(data[index]) - 1
        v = int(data[index + 1]) - 1
        edges.append((u, v))
        index += 2
    test_cases.append((N, apple_nodes, edges))

minimum_seconds_to_collect_apples(T, test_cases)

Information

Submit By
Type
Pretest
Problem
P1078 Apple on Tree
Language
Python 3 (Python 3.12.3)
Submit At
2024-08-16 16:22:21
Judged At
2024-08-16 16:22:21
Judged By
Score
0
Total Time
0ms
Peak Memory
0 Bytes