Degrees of separation

DoS1.py

from collections import defaultdict

class Graph:
    def __init__(self):
        self.graph = defaultdict(list)

    def add_edge(self, node, neighbour):  
        self.graph[node].append(neighbour)
        self.graph[neighbour].append(node)

    def zero_degree_nodes(self, node):
        return self.graph[node]

    def first_degree_nodes(self, node):
        first_degree_nodes = set()
        zero_degree_nodes = self.zero_degree_nodes(node)
        for zero_degree_node in zero_degree_nodes:
            first_degree_nodes.update(self.zero_degree_nodes(zero_degree_node))
        first_degree_nodes.difference_update(zero_degree_nodes)  # remove first-degree nodes
        first_degree_nodes.discard(node)  # remove the node itself if present
        return list(first_degree_nodes)


g = Graph()
g.add_edge('A', 'B')
g.add_edge('A', 'C')
g.add_edge('B', 'D')
g.add_edge('B', 'E')
g.add_edge('C', 'F')
g.add_edge('F', 'G')

print("Zero degree nodes of A:", g.zero_degree_nodes('A'))  # ['B', 'C']
print("First degree nodes of A:", g.first_degree_nodes('A'))  # ['D', 'E', 'F']