-
Notifications
You must be signed in to change notification settings - Fork 59
/
social_graph.py
68 lines (48 loc) · 1.64 KB
/
social_graph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
from __future__ import annotations
from collections import deque
from dataclasses import dataclass
from typing import List, Dict, Tuple
NO_FRIENDSHIP = -1
@dataclass
class User:
user_id: str
friends: List[User]
def __eq__(self: User, other: User) -> bool:
return self.user_id == other.user_id
def __hash__(self: User) -> int:
return hash(self.user_id)
def __str__(self: User) -> str:
return self.user_id
def make_friendship_graph(
nodes: List[str], edges: List[Tuple[str, str]]
) -> Dict[str, User]:
"""
We assume that the graph is connected for this example.
"""
user_graph: Dict[str, User] = {
node: User(user_id=node, friends=list()) for node in nodes
}
for edge in edges:
node_a, node_b = edge
user_graph[node_a].friends.append(user_graph[node_b])
user_graph[node_b].friends.append(user_graph[node_a])
return user_graph
def smallest_friendships(from_user: User, to_user: User) -> int:
q = deque([(from_user, 0)])
visited = {from_user}
while len(q) > 0:
top, dist = q.popleft()
if top == to_user:
return dist - 1
for friend in top.friends:
if friend not in visited:
visited.add(friend)
q.append((friend, dist + 1))
return NO_FRIENDSHIP
graph = make_friendship_graph(
nodes=["a", "b", "c", "d", "e"],
edges=[("a", "b"), ("b", "d"), ("a", "c"), ("b", "d"), ("d", "e")],
)
assert smallest_friendships(graph["a"], graph["e"]) == 2
assert smallest_friendships(graph["a"], graph["c"]) == 0
assert smallest_friendships(graph["c"], graph["e"]) == 3