K-Means Animated

An animated exploration of k-means clustering algorithm.

Drag the slider to change the number of clusters. Click the Reseed button to change the points. Click the Replay button to replay the animation. Use Step Through to manually advance the animation one step at a time.

Convergence: X's indicate the centroid (average position) for each cluster of points. Each cycle, points are reassigned to the nearest cluster and the X moves to reflect the new average position. Rings indicate the previous color value for any changed points.

Note: For a simpler example of the k-means code, see the non-animated version.

Reference: https://towardsmachinelearning.org/k-means/

Source Code for this Sketch

let points = []
let kmeans
let ksize
let autoplay = true

class Point {
	static random() {
		return new Point(random(20, width-20), random(20, height-20))
	}

	constructor(x, y) {
		this.x = x
		this.y = y
		this.k = 0
		this.changed = true
	}
}

// cross draws an x shape at the indicated x,y coordinates, with a given hue
function cross(x, y, hue) {
	const size = 5
	strokeWeight(2)
	normalizeStroke(hue, 80, 40)
	line(x-size, y-size, x+size, y+size)
	line(x-size, y+size, x+size, y-size)
}

class KMeans {
	constructor(points, k) {
		// Guard in case k is larger than the data. We can't calculate k-means
		// in that case.
		if (k > points.length) {
			return
		}

		// Initialize a random cluster for each point
		for (let i = 0; i < points.length; i++) {
			points[i].k = Math.floor(random(0, k))
		}

		let clusters = []

		// Pick random centroids and initialize each cluster (1 per k)
		for (let i = 0; i < k; i++) {
			clusters[i] = {}
			clusters[i].x = points[i].x
			clusters[i].y = points[i].y
			clusters[i].hue = i * (360 / k)
			clusters[i].size = 0
			clusters[i].totalX = 0
			clusters[i].totalY = 0
		}

		this.k = k
		this.clusters = clusters
	}

	increment() {
		let clusters = this.clusters
		let changed = 0

		// Calculate clusters
		for (let i = 0; i < points.length; i++) {
			// Calculate the nearest cluster for this point
			let lowestDistance = Number.MAX_SAFE_INTEGER
			let k = points[i].k
			for (let j = 0; j < clusters.length; j++) {
				let distance = dist(clusters[j].x, clusters[j].y, points[i].x, points[i].y)
				if (distance < lowestDistance) {
					lowestDistance = distance
					k = j
				}
			}
			if(points[i].k !== k) {
				points[i].previous = points[i].k
				points[i].changed = true
				points[i].k = k
			} else {
				points[i].previous = k
			}


			// Re-calculate centroid for the cluster
			clusters[points[i].k].size++
			clusters[points[i].k].totalX += points[i].x
			clusters[points[i].k].totalY += points[i].y
		}

		for (let i = 0; i < clusters.length; i++) {
			// Check whether the centroid has changed, and update it if it has
			let changedX = clusters[i].totalX / clusters[i].size
			if(clusters[i].x !== changedX) {
				clusters[i].x = changedX
				changed++
			}
			let changedY = clusters[i].totalY / clusters[i].size
			if (clusters[i].y !== changedY) {
				clusters[i].y = changedY
				changed++
			}

			// Zero out the counters for centroid averages
			clusters[i].size = 0
			clusters[i].totalX = 0
			clusters[i].totalY = 0
		}

		return changed === 0
	}
}

function setup() {
	randomSeed(1)
	createCanvas(600, 600, P2D, canvas)
	chrisDefaults()

	for (let i = 0; i < 100; i++) {
		points.push(Point.random())
	}

	ksize = new NativeSlider("clusters", 2, 10, 1, 4, (event, input) => {
		kmeans = new KMeans(points, input.value)
	}, null)


	new NativeButton("Randomize Points", (event) => {
		autoplay = true
		points = []
		for (let i = 0; i < 100; i++) {
			points.push(Point.random())
		}
		kmeans = new KMeans(points, ksize.value)
	})

	new NativeButton("Replay Animation", (event) => {
		autoplay = true
		kmeans = new KMeans(points, ksize.value)
	})

	new NativeButton("Step Through Animation", (event) => {
		if(autoplay) {
			autoplay = false
			kmeans = new KMeans(points, ksize.value)
		} else {
			if (kmeans.increment()) {
				autoplay = true
			}
		}
	})

	kmeans = new KMeans(points, ksize.value)
}

function normalizeStroke(hue, sat, light) {
	if(hue >= 240 || hue === 0) {
		light+=15
	}
	stroke(hue, sat, light)
}

function draw() {
	background(50, 70, 92)

	for(let i = 0; i < kmeans.clusters.length; i++) {
		cross(kmeans.clusters[i].x, kmeans.clusters[i].y, kmeans.clusters[i].hue)
	}

	for(let i = 0; i < points.length; i++) {
		if (points[i].k !== points[i].previous) {
			// Draw point in the new color
			strokeWeight(4)
			normalizeStroke(kmeans.clusters[points[i].k].hue, 80, 40)
			point(points[i].x, points[i].y)

			// Draw outline in the old color
			noFill()
			strokeWeight(1)
			normalizeStroke(points[i].previous * 360 / kmeans.clusters.length, 80, 40, .5)


			circle(points[i].x, points[i].y, 8)
			points[i].changed = false
		} else {
			// Draw point in the cluster color
			strokeWeight(4)
			normalizeStroke(kmeans.clusters[points[i].k].hue, 80, 40)
			point(points[i].x, points[i].y)
		}
	}

	if (autoplay) {
		kmeans.increment()
	}

	frameRate(2)
}

You may use all or part of the p5 code seen on this page for your own creations, without restriction. Attribution is appreciated but not required.

This sketch also includes the following sources: p5common.js . Sources are available under their respective licenses.

See all p5.js sketches