/ SeriousOJ /

Record Detail

Accepted


  

Code

def minimum_seconds_to_collect_apples(T, test_cases):
    from sys import setrecursionlimit, stdin, stdout
    from collections import deque, defaultdict

    setrecursionlimit(300000)
    output = []

    def dfs(node, parent, graph, visited, apples, distances):
        visited[node] = True
        max_dist = 0
        max_node = node
        
        for neighbor in graph[node]:
            if not visited[neighbor]:
                dist = dfs(neighbor, node, graph, visited, apples, distances)
                if dist > max_dist:
                    max_dist = dist
                    max_node = neighbor
        
        distances[node] = max_dist
        return max_dist + 1

    for case in test_cases:
        N, apple_nodes, edges = case
        if not any(apple_nodes):
            output.append("0")
            continue
        
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        apples = [i for i in range(N) if apple_nodes[i]]
        
        # Start DFS from the first apple node to find the farthest apple node
        visited = [False] * N
        distances = [0] * N
        max_node = apples[0]
        dfs(max_node, -1, graph, visited, apple_nodes, distances)
        
        # Now find the farthest node from the previous farthest node
        visited = [False] * N
        distances = [0] * N
        max_dist = dfs(max_node, -1, graph, visited, apple_nodes, distances)
        
        output.append(str(max_dist - 1))
    
    stdout.write("\n".join(output) + "\n")

# Sample input processing
import sys
input = sys.stdin.read
data = input().strip().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:21:28
Judged At
2024-08-16 16:21:28
Judged By
Score
0
Total Time
0ms
Peak Memory
0 Bytes