recommander.py

Recommander.py

'''
Use inverted index algorithm to enable simple movie recommandation
'''
import collections


class User:
    def __init__(self, uid):
        self.uid = uid
        self.movies = []

    def add_movies(self, movies):
        self.movies.extend(movies)

    def __str__(self):
        return 'User<%d>' % self.uid

class Movie:
    def __init__(self, mid):
        self.mid = mid

    def __str():
        return 'Movie<%d>' % self.mid

class Recommander:
    @staticmethod
    def find_similar_user(uid, users):
        inverted_index = collections.defaultdict(list)
        for muid in users:
            user = users[muid]
            for movie in user.movies:
                inverted_index[movie.mid].append(user.uid)
        # caculate similarity
        susers = collections.defaultdict(int)
        for movie in users[uid].movies:
            for suid in inverted_index[movie.mid]:
                if suid != uid:
                    susers[suid] += 1
        smax = max(susers.values())
        for suid, val in susers.iteritems():
            if val == smax:
                return users[suid]

if __name__ == '__main__':
    users = {u: User(u) for u in range(1, 4)}
    movies = {m: Movie(m) for m in [1, 3, 4, 5, 7, 9]}
    users[1].add_movies([movies[1], movies[3], movies[7]])
    users[2].add_movies([movies[3], movies[5]])
    users[3].add_movies([movies[1], movies[3], movies[4], movies[9]])

    print('User<1> and %s have similar movie insterest.' % Recommander.find_similar_user(1, users))