Cod sursa(job #3258976)

Utilizator MateitzaMatei Dinu Mateitza Data 24 noiembrie 2024 16:25:36
Problema Diametrul unui arbore Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 1.86 kb
class Solution:
    def solveProblem(self):
        def read_file():
            graph = {}
            with open('darb.in', 'r') as f:
                n = int(f.readline())
                for _ in range(n - 1):
                    u, v = map(int, f.readline().split())
                    if u not in graph: graph[u] = []
                    if v not in graph: graph[v] = []
                    graph[u].append(v)
                    graph[v].append(u)
            return graph
        
        def bfs(start, graph):
            # Queue will store (node, distance) pairs
            queue = [(start, 0)]
            visited = {start}
            max_dist = 0
            farthest_node = start
            
            idx = 0
            # Using list as queue and index for better memory efficiency
            while idx < len(queue):
                node, dist = queue[idx]
                idx += 1
                
                # Update farthest node if needed
                if dist > max_dist:
                    max_dist = dist
                    farthest_node = node
                
                # Process neighbors
                for neighbor in graph[node]:
                    if neighbor not in visited:
                        visited.add(neighbor)
                        queue.append((neighbor, dist + 1))
            
            return max_dist, farthest_node

        # Build graph
        tree = read_file()
        
        # Find farthest node from any starting point
        start = next(iter(tree))
        _, far_node = bfs(start, tree)
        
        # Find the diameter starting from the farthest node
        diameter, _ = bfs(far_node, tree)
        
        # Write result
        with open('darb.out', 'w') as f:
            f.write(str(diameter + 1))

if __name__ == "__main__":
    Solution().solveProblem()