from collections import deque my_input = open("inputs/input_4", "r").read().split("\n") start = "STT" class City: def __init__(self, name): self.name = name self.connects = set() cities = {} for line in my_input: city_left, city_right = line.split(" <-> ") if city_left not in cities: cities[city_left] = City(city_left) if city_right not in cities: cities[city_right] = City(city_right) cities[city_left].connects.add(cities[city_right]) cities[city_right].connects.add(cities[city_left]) p1 = len(cities) p2 = 0 p3 = 0 q = deque([(0, cities[start])]) seen = set() while q: dist, city = q.popleft() if city in seen: continue seen.add(city) p3 += dist if dist <= 3: p2 += 1 for sub_city in city.connects: q.append((dist + 1, sub_city)) print(p1) print(p2) print(p3)