Intial commit, with some work done to the 'shortest_path' method.
This commit is contained in:
@@ -0,0 +1 @@
|
|||||||
|
__pycache__
|
||||||
+172
@@ -0,0 +1,172 @@
|
|||||||
|
import csv
|
||||||
|
import sys
|
||||||
|
|
||||||
|
from util import Node, StackFrontier, QueueFrontier
|
||||||
|
|
||||||
|
# Maps names to a set of corresponding person_ids
|
||||||
|
names = {}
|
||||||
|
|
||||||
|
# Maps person_ids to a dictionary of: name, birth, movies (a set of movie_ids)
|
||||||
|
people = {}
|
||||||
|
|
||||||
|
# Maps movie_ids to a dictionary of: title, year, stars (a set of person_ids)
|
||||||
|
movies = {}
|
||||||
|
|
||||||
|
|
||||||
|
def load_data(directory):
|
||||||
|
"""
|
||||||
|
Load data from CSV files into memory.
|
||||||
|
"""
|
||||||
|
# Load people
|
||||||
|
with open(f"{directory}/people.csv", encoding="utf-8") as f:
|
||||||
|
reader = csv.DictReader(f)
|
||||||
|
for row in reader:
|
||||||
|
people[row["id"]] = {
|
||||||
|
"name": row["name"],
|
||||||
|
"birth": row["birth"],
|
||||||
|
"movies": set()
|
||||||
|
}
|
||||||
|
if row["name"].lower() not in names:
|
||||||
|
names[row["name"].lower()] = {row["id"]}
|
||||||
|
else:
|
||||||
|
names[row["name"].lower()].add(row["id"])
|
||||||
|
|
||||||
|
# Load movies
|
||||||
|
with open(f"{directory}/movies.csv", encoding="utf-8") as f:
|
||||||
|
reader = csv.DictReader(f)
|
||||||
|
for row in reader:
|
||||||
|
movies[row["id"]] = {
|
||||||
|
"title": row["title"],
|
||||||
|
"year": row["year"],
|
||||||
|
"stars": set()
|
||||||
|
}
|
||||||
|
|
||||||
|
# Load stars
|
||||||
|
with open(f"{directory}/stars.csv", encoding="utf-8") as f:
|
||||||
|
reader = csv.DictReader(f)
|
||||||
|
for row in reader:
|
||||||
|
try:
|
||||||
|
people[row["person_id"]]["movies"].add(row["movie_id"])
|
||||||
|
movies[row["movie_id"]]["stars"].add(row["person_id"])
|
||||||
|
except KeyError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
if len(sys.argv) > 2:
|
||||||
|
sys.exit("Usage: python degrees.py [directory]")
|
||||||
|
directory = sys.argv[1] if len(sys.argv) == 2 else "large"
|
||||||
|
|
||||||
|
# Load data from files into memory
|
||||||
|
print("Loading data...")
|
||||||
|
load_data(directory)
|
||||||
|
print("Data loaded.")
|
||||||
|
|
||||||
|
source = person_id_for_name(input("Name: "))
|
||||||
|
if source is None:
|
||||||
|
sys.exit("Person not found.")
|
||||||
|
target = person_id_for_name(input("Name: "))
|
||||||
|
if target is None:
|
||||||
|
sys.exit("Person not found.")
|
||||||
|
|
||||||
|
path = shortest_path(source, target)
|
||||||
|
|
||||||
|
if path is None:
|
||||||
|
print("Not connected.")
|
||||||
|
else:
|
||||||
|
degrees = len(path)
|
||||||
|
print(f"{degrees} degrees of separation.")
|
||||||
|
path = [(None, source)] + path
|
||||||
|
for i in range(degrees):
|
||||||
|
person1 = people[path[i][1]]["name"]
|
||||||
|
person2 = people[path[i + 1][1]]["name"]
|
||||||
|
movie = movies[path[i + 1][0]]["title"]
|
||||||
|
print(f"{i + 1}: {person1} and {person2} starred in {movie}")
|
||||||
|
|
||||||
|
|
||||||
|
def shortest_path(source, target):
|
||||||
|
"""
|
||||||
|
Returns the shortest list of (movie_id, person_id) pairs
|
||||||
|
that connect the source to the target.
|
||||||
|
|
||||||
|
If no possible path, returns None.
|
||||||
|
"""
|
||||||
|
frontier = QueueFrontier()
|
||||||
|
explored = StackFrontier()
|
||||||
|
frontier.add(Node(source, None, None))
|
||||||
|
results = [[]]
|
||||||
|
|
||||||
|
while frontier:
|
||||||
|
currentNode = frontier.remove()
|
||||||
|
|
||||||
|
if currentNode.state == target:
|
||||||
|
tmp = currentNode
|
||||||
|
print("Found %s" %(people[tmp.state]))
|
||||||
|
print("Found!")
|
||||||
|
cost = 0
|
||||||
|
r = []
|
||||||
|
r.append(tmp)
|
||||||
|
index = len(results)
|
||||||
|
|
||||||
|
while tmp:
|
||||||
|
if not tmp.parent:
|
||||||
|
r.append(tmp.parent)
|
||||||
|
print("Parent found! %s" %(people[tmp.state]))
|
||||||
|
|
||||||
|
tmp = tmp.parent
|
||||||
|
cost+=1
|
||||||
|
results.append(r)
|
||||||
|
print("Estimated Cost %d" %(cost))
|
||||||
|
#break
|
||||||
|
|
||||||
|
if not explored.contains_state(currentNode.state):
|
||||||
|
# currentNode.state has person ID.
|
||||||
|
neighbors = neighbors_for_person(currentNode.state)
|
||||||
|
explored.add(currentNode)
|
||||||
|
|
||||||
|
for i in neighbors:
|
||||||
|
frontier.add(Node(i[1], currentNode, None))
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
def person_id_for_name(name):
|
||||||
|
"""
|
||||||
|
Returns the IMDB id for a person's name,
|
||||||
|
resolving ambiguities as needed.
|
||||||
|
"""
|
||||||
|
person_ids = list(names.get(name.lower(), set()))
|
||||||
|
if len(person_ids) == 0:
|
||||||
|
return None
|
||||||
|
elif len(person_ids) > 1:
|
||||||
|
print(f"Which '{name}'?")
|
||||||
|
for person_id in person_ids:
|
||||||
|
person = people[person_id]
|
||||||
|
name = person["name"]
|
||||||
|
birth = person["birth"]
|
||||||
|
print(f"ID: {person_id}, Name: {name}, Birth: {birth}")
|
||||||
|
try:
|
||||||
|
person_id = input("Intended Person ID: ")
|
||||||
|
if person_id in person_ids:
|
||||||
|
return person_id
|
||||||
|
except ValueError:
|
||||||
|
pass
|
||||||
|
return None
|
||||||
|
else:
|
||||||
|
return person_ids[0]
|
||||||
|
|
||||||
|
|
||||||
|
def neighbors_for_person(person_id):
|
||||||
|
"""
|
||||||
|
Returns (movie_id, person_id) pairs for people
|
||||||
|
who starred with a given person.
|
||||||
|
"""
|
||||||
|
movie_ids = people[person_id]["movies"]
|
||||||
|
neighbors = set()
|
||||||
|
for movie_id in movie_ids:
|
||||||
|
for person_id in movies[movie_id]["stars"]:
|
||||||
|
neighbors.add((movie_id, person_id))
|
||||||
|
return neighbors
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
+344277
File diff suppressed because it is too large
Load Diff
+1044500
File diff suppressed because it is too large
Load Diff
+1189595
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,6 @@
|
|||||||
|
id,title,year
|
||||||
|
112384,"Apollo 13",1995
|
||||||
|
104257,"A Few Good Men",1992
|
||||||
|
109830,"Forrest Gump",1994
|
||||||
|
93779,"The Princess Bride",1987
|
||||||
|
95953,"Rain Man",1988
|
||||||
|
@@ -0,0 +1,17 @@
|
|||||||
|
id,name,birth
|
||||||
|
102,"Kevin Bacon",1958
|
||||||
|
129,"Tom Cruise",1962
|
||||||
|
144,"Cary Elwes",1962
|
||||||
|
158,"Tom Hanks",1956
|
||||||
|
1597,"Mandy Patinkin",1952
|
||||||
|
163,"Dustin Hoffman",1937
|
||||||
|
1697,"Chris Sarandon",1942
|
||||||
|
193,"Demi Moore",1962
|
||||||
|
197,"Jack Nicholson",1937
|
||||||
|
200,"Bill Paxton",1955
|
||||||
|
398,"Sally Field",1946
|
||||||
|
420,"Valeria Golino",1965
|
||||||
|
596520,"Gerald R. Molen",1935
|
||||||
|
641,"Gary Sinise",1955
|
||||||
|
705,"Robin Wright",1966
|
||||||
|
914612,"Emma Watson",1990
|
||||||
|
@@ -0,0 +1,21 @@
|
|||||||
|
person_id,movie_id
|
||||||
|
102,104257
|
||||||
|
102,112384
|
||||||
|
129,104257
|
||||||
|
129,95953
|
||||||
|
144,93779
|
||||||
|
158,109830
|
||||||
|
158,112384
|
||||||
|
1597,93779
|
||||||
|
163,95953
|
||||||
|
1697,93779
|
||||||
|
193,104257
|
||||||
|
197,104257
|
||||||
|
200,112384
|
||||||
|
398,109830
|
||||||
|
420,95953
|
||||||
|
596520,95953
|
||||||
|
641,109830
|
||||||
|
641,112384
|
||||||
|
705,109830
|
||||||
|
705,93779
|
||||||
|
@@ -0,0 +1,38 @@
|
|||||||
|
class Node():
|
||||||
|
def __init__(self, state, parent, action):
|
||||||
|
self.state = state
|
||||||
|
self.parent = parent
|
||||||
|
self.action = action
|
||||||
|
|
||||||
|
|
||||||
|
class StackFrontier():
|
||||||
|
def __init__(self):
|
||||||
|
self.frontier = []
|
||||||
|
|
||||||
|
def add(self, node):
|
||||||
|
self.frontier.append(node)
|
||||||
|
|
||||||
|
def contains_state(self, state):
|
||||||
|
return any(node.state == state for node in self.frontier)
|
||||||
|
|
||||||
|
def empty(self):
|
||||||
|
return len(self.frontier) == 0
|
||||||
|
|
||||||
|
def remove(self):
|
||||||
|
if self.empty():
|
||||||
|
raise Exception("empty frontier")
|
||||||
|
else:
|
||||||
|
node = self.frontier[-1]
|
||||||
|
self.frontier = self.frontier[:-1]
|
||||||
|
return node
|
||||||
|
|
||||||
|
|
||||||
|
class QueueFrontier(StackFrontier):
|
||||||
|
|
||||||
|
def remove(self):
|
||||||
|
if self.empty():
|
||||||
|
raise Exception("empty frontier")
|
||||||
|
else:
|
||||||
|
node = self.frontier[0]
|
||||||
|
self.frontier = self.frontier[1:]
|
||||||
|
return node
|
||||||
Reference in New Issue
Block a user