scikit-learnを使ってMovieLensの映画評価データをクラスタリングしてみた

2014-11-07
このエントリーをはてなブックマークに追加

使ったのはMovieLens 100k
「100,000 ratings from 1000 users on 1700 movies.」
だそうです。

from sklearn.cluster import KMeans


def main():

    _items = []
    _f = open('./rating.csv' )
    _lines = _f.readlines()

    for _line in _lines:
        _items.append(_line.split(','))

    _f.close()

    km = KMeans(init='k-means++')
    km.fit(_items)

    print(km.labels_)


if __name__ == '__main__':
    main()

すげえ簡単。。。

CSVは1行が1人分で全部の映画のレーティング情報(0が観てない、評価は1から5)が入っています。1000ユーザーが1700本の映画を評価(未鑑賞含む)してます。

だいたいこんな感じ。

5,3,4,3,3,5,4,1,5,3,2,5,5,5,5,5,3,4,,,,,,,,
4,0,0,0,0,0,0,0,0,2,0,0,4,4,0,0,0,0,,,,,,,,

得られる結果は

[0 1 4 4 3 2 6 3 4 6 3 7 6 2 1 5 1 6 4 7 1 3 3 7,,,,,,,

インデックス2番のユーザーさんと3番のユーザーさんは好きな映画が似る傾向にありそうですね。

Python 3.4.2、scikit-learn 0.15.2 で試しました。