system-design-primer/solutions/system_design/social_graph/social_graph_snippets.py

# -*- coding: utf-8 -*-
from collections import deque
from enum import Enum


class State(Enum) :
    unvisited = 0
    visited = 1


class Graph(object) :

    def bfs(self, source, dest) :
        if source is None:
            return False
        queue = deque() 
        queue.append(source) 
        source.visit_state = State.visited
        while queue:
            node = queue.popleft() 
            print(node) 
            if dest is node:
                return True
            for adjacent_node in node.adj_nodes.values() :
                if adjacent_node.visit_state == State.unvisited:
                    queue.append(adjacent_node) 
                    adjacent_node.visit_state = State.visited
        return False


class Person(object) :

    def __init__(self, id, name) :
        self.id = id
        self.name = name
        self.friend_ids = []


class LookupService(object) :

    def __init__(self) :
        self.lookup = {}  # key: person_id, value: person_server

    def get_person(self, person_id) :
        person_server = self.lookup[person_id]
        return person_server.people[person_id]


class PersonServer(object) :

    def __init__(self) :
        self.people = {}  # key: person_id, value: person

    def get_people(self, ids) :
        results = []
        for id in ids:
            if id in self.people:
                results.append(self.people[id]) 
        return results


class UserGraphService(object) :

    def __init__(self, person_ids, lookup) :
        self.lookup = lookup
        self.person_ids = person_ids
        self.visited_ids = set() 

    def bfs(self, source, dest) :
        # Use self.visited_ids to track visited nodes
        # Use self.lookup to translate a person_id to a Person
        pass