K-Means

An exploration of k-means clustering algorithm.

Drag the slider to change the number of clusters. Click the Reseed button to change the points.

See also: the animated version

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

Source Code for this Sketch

let points = []
let clusters = []
let ksize

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
	}
}

// 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)
	stroke(hue, 80, 40)
	line(x-size, y-size, x+size, y+size)
	line(x-size, y+size, x+size, y-size)
}

function kMeans(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))
	}

	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
	}

	let foundChanges = 1
	let totalChanges = 0
	let cycleLimit = 30

	// Calculate clusters
	while (foundChanges !== 0) {
		foundChanges = 0

		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
				}
			}
			points[i].k = 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
				foundChanges++
			}
			let changedY = clusters[i].totalY / clusters[i].size
			if (clusters[i].y !== changedY) {
				clusters[i].y = changedY
				foundChanges++
			}

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

		totalChanges++

		// Guard against runaway loops
		cycleLimit--
		if(cycleLimit <= 0) {
			console.error("Reached cycle limit")
			break
		}
	}

	console.log(k, "clusters identified in", totalChanges, "passes")
}

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

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

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

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

	frameRate(12)
	kMeans(ksize.value)
}

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

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

	strokeWeight(4)
	for(let i = 0; i < points.length; i++) {
		stroke(clusters[points[i].k].hue, 80, 40)
		point(points[i].x, points[i].y)
	}
}

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