RでEMアルゴリズムとしてのK-meansクラスタリング

PRML 9.1に記載の通り、K-meansクラスタリングが行われる過程と、コスト関数の収束の様子を示します。

library(MASS)
frame()
set.seed(0)
par(mfrow=c(4, 4))
par(mar=c(2.5, 2.5, 1, 0.1))
par(mgp=c(1.3, 0.5, 0))
xrange <- c(-2, 2)
yrange <- c(-2, 2)
D <- 2
K <- 2
data(faithful)
x <- as.matrix(faithful)
N <- nrow(x)
#N <- 100
#x <- mvrnorm(N / 2, c(-1, -1), matrix(c(.4, 0, 0, .4), D))
#x <- rbind(x, mvrnorm(N / 2, c(1, 1), matrix(c(.4, 0, 0, .4), D)))
x <- t((t(x) - apply(x, 2, mean)) / apply(x, 2, sd))  # normalize
mu <- matrix(c(-1, 1, 1, -1), K, byrow=T)
r <- matrix(NA, nrow=N, ncol=K)
j <- numeric()

cost <- function() {
    sum(sapply(1:N, function(n)
        sum(r[n, ] * rowSums(t(x[n, ] - t(mu)) ^ 2))
        ))
}

iteration <- 0
repeat {
    iteration <- iteration + 1

    # E step
    for (n in 1:N) {
        diff <- numeric(K)
        for (k in 1:K) {
            diff[k] <- sum((x[n, ] - mu[k, ]) ^ 2)
        }
        minK <- which.min(diff)
        r[n, minK] <- 1
        r[n, -minK] <- 0
    }
    j <- c(j, cost())
    plot(x, xlim=xrange, ylim=yrange, col=hsv(0.4 * apply(r, 1, function(r) which(r == 1) - 1)), pch=16)
    par(new=T)
    plot(mu, xlim=xrange, ylim=yrange, xlab="", ylab="", col=hsv((1:K - 1) * 0.4, 1, 0.8), pch=0, cex=2)
    title(paste0("E step#", iteration))

    # M step
    for (k in 1:K) {
        mu[k, ] <- colSums(x * r[, k]) / sum(r[, k])
    }
    j <- c(j, cost())
    plot(x, xlim=xrange, ylim=yrange, col=hsv(0.4 * apply(r, 1, function(r) which(r == 1) - 1)), pch=1)
    par(new=T)
    plot(mu, xlim=xrange, ylim=yrange, xlab="", ylab="", col=hsv((1:K - 1) * 0.4, 1, 0.8), pch=15, cex=2)
    title(paste0("M step#", iteration))

    if (length(j) > 2 && j[length(j)] == j[length(j) - 2]) {
        break
    }
}

names(j) <- paste0(rep(1:(length(j) / 2), each=2), c("E", "M"))
plot(j, type="o", axes=F, xlab="iteration")
axis(1, at=1:length(j), labels=names(j))
axis(2)
title("cost")