32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
|
from collections import deque
|
||
|
|
|
||
|
|
from app.models.evolution import Evolution
|
||
|
|
|
||
|
|
|
||
|
|
def build_families(evolutions: list[Evolution]) -> dict[int, list[int]]:
|
||
|
|
"""Build pokemon_id -> family members mapping using BFS on evolution graph."""
|
||
|
|
adj: dict[int, set[int]] = {}
|
||
|
|
for evo in evolutions:
|
||
|
|
adj.setdefault(evo.from_pokemon_id, set()).add(evo.to_pokemon_id)
|
||
|
|
adj.setdefault(evo.to_pokemon_id, set()).add(evo.from_pokemon_id)
|
||
|
|
|
||
|
|
visited: set[int] = set()
|
||
|
|
pokemon_to_family: dict[int, list[int]] = {}
|
||
|
|
for node in adj:
|
||
|
|
if node in visited:
|
||
|
|
continue
|
||
|
|
component: list[int] = []
|
||
|
|
queue = deque([node])
|
||
|
|
while queue:
|
||
|
|
current = queue.popleft()
|
||
|
|
if current in visited:
|
||
|
|
continue
|
||
|
|
visited.add(current)
|
||
|
|
component.append(current)
|
||
|
|
for neighbor in adj.get(current, set()):
|
||
|
|
if neighbor not in visited:
|
||
|
|
queue.append(neighbor)
|
||
|
|
for member in component:
|
||
|
|
pokemon_to_family[member] = component
|
||
|
|
return pokemon_to_family
|