Cod sursa(job #3258972)

Utilizator MateitzaMatei Dinu Mateitza Data 24 noiembrie 2024 16:15:47
Problema Diametrul unui arbore Scor 50
Compilator py Status done
Runda Arhiva educationala Marime 2.61 kb
class Solution:
    def solveProblem(self):
        hashmap_node_neighbors = {}
        
        def read_file(file_name: str):
            with open(file_name, 'r') as f:
                lines = f.readlines()
                
            n = int(lines[0])
            for i in range(1, n):
                node, neighbor = map(int, lines[i].split())
                
                if node not in hashmap_node_neighbors:
                    hashmap_node_neighbors[node] = []
                if neighbor not in hashmap_node_neighbors:
                    hashmap_node_neighbors[neighbor] = []
                
                hashmap_node_neighbors[node].append(neighbor)
                hashmap_node_neighbors[neighbor].append(node)
            
            return lines, hashmap_node_neighbors
        
        def dfs(node: int, depth: int, visited: set):  # Added visited set parameter
            # Mark current node as visited
            visited.add(node)
            
            # If the current node has no unvisited neighbors, return current depth
            unvisited_neighbors = [n for n in hashmap_node_neighbors[node] if n not in visited]
            if not unvisited_neighbors:
                return (depth, node)
            
            maxval = (depth, node)
            for nextVal in unvisited_neighbors:  # Only visit unvisited neighbors
                dfirst, node_d = dfs(nextVal, depth+1, visited)
                if dfirst > maxval[0]:
                    maxval = (dfirst, node_d)
            return maxval
        
        def dfs2(node: int, depth: int, visited: set):  # Changed visited to set
            visited.add(node)
            
            unvisited = [n for n in hashmap_node_neighbors[node] if n not in visited]
            if not unvisited:
                return (depth, node)
            
            maxval = (depth, node)
            for nextVal in unvisited:
                dfirst, node_d = dfs2(nextVal, depth+1, visited)
                if dfirst > maxval[0]:
                    maxval = (dfirst, node_d)
            return maxval

        # Main execution
        filelines, lines = read_file('darb.in')
        
        # First DFS to find one end of the diameter
        start_node = list(lines.keys())[0]
        dmaxim, node = dfs(start_node, 1, set())  # Added empty set for visited nodes
        
        # Second DFS from the found node
        dmaxim_2, node2 = dfs2(node, 1, set())
        
        # Write result to output file
        with open('darb.out', 'w') as f:
            f.write(str(dmaxim_2))

sol = Solution()
sol.solveProblem()