K-meansで遊んでみた

研究室の輪講でK-meansを実装する課題を出すために、自分でも実装してみました。

from copy import deepcopy
from random import shuffle

def kmeans(rows, k=4, distance=None):
    if not distance:
        distance = euclidean

    dim = len(rows[0])  # 特徴空間の次元数

    # データの中からcentroidをk個ランダムに選択
    shuffle(rows)
    centroids = rows[:k]

    result = []

    prev = None
    while True:
        clusters = [[] for i in xrange(k)]
        
        # それぞれのデータに対し、最も近い重心を探す
        for j, row in enumerate(rows):
            ds = map(lambda x: distance(row, x), centroids)
            m = min(ds)
            clusters[ds.index(m)].append(j)

        # 結果が前回と同じ場合は終了
        if clusters == prev: break
        prev = clusters

        # 重心をメンバーの平均に移動する
        for j, member in enumerate(clusters):
            centroid = [0.0 for i in xrange(dim)]
            for m in member:
                row = rows[m]
                for n, val in enumerate(row):
                    centroid[n] += val
            for m in xrange(dim):
                centroid[m] /= len(member)
            centroids[j] = centroid

        # 可視化するために途中経過も結果に含める
        result.append((clusters, deepcopy(centroids)))

    return result

def euclidean(v1, v2):
    'ユークリッド距離'
    dim = len(v1)  # 特徴ベクトルの次元数
    assert len(v1) == len(v2)

    sum_sq = 0.0
    for i in xrange(dim):
        sum_sq += pow(v1[i] - v2[i], 2)

    return 1 - 1 / max(pow(sum_sq, 0.5), 0.00001)

これに対して、次のようなコードを実行してみると

from random import randint

def generate_data(size_x=10, size_y=10, t=4):
    data = []
    for i in xrange(size_y):
        line = ""
        for j in xrange(size_x):
            if randint(0, 10) < t:
                line += "*"
            else:
                line += " "
        data.append(line)
    print 'data set'
    print "\n".join(data)

    rows = []
    for y_, line in enumerate(data):
        y = len(data) - y_ - 1
        for x, v in enumerate(line):
            if v == "*":
                rows.append((x, y))
    return rows

def print_result(size_x, size_y, cluster, centroid):
    rounded_centroid = [(int(round(x1)), int(round(x2))) for x1, x2 in centroid]

    print rounded_centroid
    for y in xrange(size_y-1, -1, -1):
        line = ""
        for x in xrange(size_x):
            for i, member in enumerate(cluster):
                if (x, y) in rounded_centroid:
                    line += "*"
                    break
                elif (x, y) in rows and rows.index((x, y)) in member:
                    line += str(i)
                    break
            else:
                line += " "
        print line

# 20*10 の大きさの平面に2割程度データを設置する
size_x = 20
size_y = 10
t = 2

# データセットの生成
rows = generate_data(size_x, size_y, t)

# kmeansの実行。今回は4クラスタに分類する
results = kmeans(rows, k=4)

print 'result'
for i, (cluster, centroid) in enumerate(results):
    print ''
    print i, 
    print_result(size_x, size_y, cluster, centroid)

例えばこういう感じに表示されます。
結果の出力の中で＊と表示されているのが、クラスタの重心で、数字は分類されているクラスタの番号を表しています。
反復するたびに少しずつ移動しているのが分かりますね。

# 入力されるデータ
data set
  * ***    *     * *
   * *  *          *
        *     *    *
   *   **  *      * 
                   *
* * *     *         
          *  * *   *
* *           *     
              *     
   * *            * 

# 1回目の反復での結果
0 [(13, 3), (2, 4), (14, 7), (3, 6)]
  3 333    2     2 2
   3 3  3          2
        3     *    2
   *   33  2      2 
                   2
1 * 1     0         
          0  * 0   0
1 1           0     
              0     
   1 1            0 

# 2回目の反復での結果
1 [(14, 2), (2, 2), (16, 7), (5, 8)]
  3 333    2     2 2
   3 *  3          2
        3     2 *  2
   3   33  0      2 
                   2
1 1 1     0         
          0  0 0   0
1 *           *     
              0     
   1 1            0 

# 3回目の反復での結果
2 [(14, 3), (2, 2), (17, 8), (5, 8)]
  3 333    3     2 2
   3 *  3        * 2
        3     2    2
   3   33  0      2 
                   2
1 1 1     0         
          0  0*0   2
1 *           0     
              0     
   1 1            0