dfs on python comparing dict keys as strings

graph_depth_first_search.py

"""Depth-first search implemented in python."""

from collections import defaultdict
from typing import DefaultDict, Dict, List, Union

Raw_data = Dict[str, Dict[str, float]]
Graph = DefaultDict[str, List[str]]
Path = Union[None, List[str]]


def construct_graph(raw_data: Raw_data) -> Graph:
    graph = defaultdict(list)  # the graph
    for pair in raw_data:  # visit every edge
        origin, vertexes = pair.split("/")  # get from and to vertexes
        graph[origin].append(vertexes)  # add this edge in our structure
    return graph  # return new graph form


def depth_first_search(graph: Graph, origin: str, distance: int) -> Path:
    if distance == 2:  # we have a 'distance' from our start
        return [origin]  # if found, return
    for vertexes in graph.get(origin, []):  # check all neighbors
        answer = depth_first_search(
            graph,
            vertexes,
            distance + 1,
        )  # run dfs in every neighbor with dist+1
        if answer:  # and if dfs found something
            answer.append(origin)  # store it the answer
            return answer  # and return it


def main(raw_data: Raw_data) -> Path:
    graph = construct_graph(raw_data)
    answer = []
    for vertexes in graph.keys():  # try to find a path
        path_found = depth_first_search(
            graph,
            vertexes,
            0,
        )  # starting with 0 distance
        if path_found:
            answer.append(path_found)
    return answer


if __name__ == "__main__":
    raw_data = {
        "EZ/TC": {"id": 1, "unit": 9.0},
        "LM/TH": {"id": 2, "unit": 8.0},
        "CD/EH": {"id": 3, "unit": 7.0},
        "EH/TC": {"id": 4, "unit": 6.0},
        "LM/TC": {"id": 5, "unit": 5.0},
        "CD/TC": {"id": 6, "unit": 4.0},
        "BT/TH": {"id": 7, "unit": 3.0},
        "BT/TX": {"id": 8, "unit": 2.0},
        "TX/TH": {"id": 9, "unit": 1.0},
    }
    print(main(raw_data))