changes to the code

solutions/system_design/social_graph/social_graph_snippets.py

@@ -1,72 +1,67 @@
-# -*- coding: utf-8 -*-
 from collections import deque
 from enum import Enum
 
 
-class State(Enum):
+class Person:
-    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 = []
 
+    def __repr__(self):
+        return f"Person({self.id}, {self.name})"
 
-class LookupService(object):
 
+class LookupService:
     def __init__(self):
-        self.lookup = {}  # key: person_id, value: person_server
+        self.lookup = {}  # key: person_id, value: PersonServer
 
     def get_person(self, person_id):
-        person_server = self.lookup[person_id]
+        person_server = self.lookup.get(person_id)
-        return person_server.people[person_id]
+        if person_server:
+            return person_server.people.get(person_id)
+        return None
 
 
-class PersonServer(object):
+class PersonServer:
-
     def __init__(self):
-        self.people = {}  # key: person_id, value: person
+        self.people = {}  # key: person_id, value: Person
 
     def get_people(self, ids):
-        results = []
+        return [self.people[id] for id in ids if id in self.people]
-        for id in ids:
-            if id in self.people:
-                results.append(self.people[id])
-        return results
 
 
-class UserGraphService(object):
+class UserGraphService:
-
+    def __init__(self, person_ids, lookup: LookupService):
-    def __init__(self, person_ids, lookup):
         self.lookup = lookup
         self.person_ids = person_ids
         self.visited_ids = set()
 
-    def bfs(self, source, dest):
+    def bfs(self, source_id, dest_id):
-        # Use self.visited_ids to track visited nodes
+        source = self.lookup.get_person(source_id)
-        # Use self.lookup to translate a person_id to a Person
+        dest = self.lookup.get_person(dest_id)
-        pass
+
+        if source is None or dest is None:
+            return False
+
+        queue = deque()
+        queue.append(source)
+        self.visited_ids.add(source.id)
+
+        while queue:
+            current_person = queue.popleft()
+            print(current_person)
+
+            if current_person.id == dest.id:
+                return True
+
+            # Get actual Person objects for friends
+            for friend_id in current_person.friend_ids:
+                if friend_id not in self.visited_ids:
+                    friend = self.lookup.get_person(friend_id)
+                    if friend:
+                        queue.append(friend)
+                        self.visited_ids.add(friend.id)
+
+        return False
+