使ったのは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 で試しました。