Clustering the US population: observation-weighted k-means | by Carl Anderson | Towards Data Science

Thanks to Chris Decker who provided the following info:
For anyone discovering this post in recent years: scikit learn implemented a ‘sample_weight’ parameter into KMeans as of 0.20.0 in 2018. No need to roll your own anymore.
— — — — — -In this post, I detail a form of k-means clustering in which weights are associated with individual observations.
k-means clustering has been a workhorse of machine learning for almost 60 years. It is simple, effective, and intuitive. One of the reasons that I love it, is that you can plot cluster assignments over time and see it learning.
However, it doesn’t scale particularly well. On a laptop, can you cluster 10,000 points? Sure, no problem. A million? Maybe, but it will be slow. A hundred million points? Fuhgeddaboutit!
Imagine that we wanted to cluster the U.S. population. Why? We could set k=48 and determine what the lower 48 states might look like based on current centers of population. Perhaps, we want to set up k=1000 distribution centers across the country. Maybe we want to set k=10 million Starbucks locations.
The U.S. population has more than 300 million people. How would you cluster them using k-means? Ignoring parallel k-means, let’s constrain it to run on a single laptop in say less than 3 minutes. Ideas?
One thing that you could do is sample the data; that is, run with a reduced, hopefully representative, subset of the data. Another approach would be to aggregate the data, drastically reducing the number of points, but associate each with the original sample size that each aggregate point represents. In other words, weight each point.
Where might we get such a dataset? The U.S. Census Bureau. They provide a very convenient dataset consisting of the number of inhabitants for each Zip code. There are about 43,000 Zip codes in the U.S, a number we can comfortably cluster on a laptop. Imagine then, we have a data file consisting of Zip code, a latitude-longitude pair (which are the x-y coordinates that k-means works on), and the number of inhabitants in that Zip (the weight):
“zip”,”state”,”latitude”,”longitude”,”population”“00601”,”PR”,18.180103,-66.74947,19143“00602”,”PR”,18.363285,-67.18024,42042...99929,”AK”,56.409507,-132.33822,242499950,”AK”,55.875767,-131.46633,47
Surely we can go to scikit-learn or R or other major machine learning library and run some weighted k-means algorithm. Unfortunately, not. There are weighted k-means in a few of those libraries but they are not the sort that we want. They provide weights not for the observations but for the features.
That is, with feature-weights we could specify that latitude should influence the centroids more than longitude. However, we want to specify observation-weights such that Zip code=10012 (a dense Manhattan Zip) has far greater draw on a centroid than 59011 (a vast, low density region in Montana).
The core idea is very intuitive. Take four equally-weighted points in the (x,y) plane and the centroid is (mean(x), mean(y)).
If we apply weights, say w=(13,7,4,4), then the point with weight 13 has far greater gravity and should draw the cluster center much closer to it. The weighted centroid is now: (weighted.mean(x,w),weighted.mean(y,w))
The distance metric can be our usual Euclidean distance. However, as we are dealing with latitude-longitude pairs, the correct as-the-crow-flies distance metric is something called the Haversine distance so we’ll use that.
I was not able to find an implementation in major library. Even stackoverflow failed me. One suggestion from there was to replicate the data points. That is, if the weighted data were as shown (top), one would created the equivalent unweighted dataset (below):
Obviously, this has the problem of scaling up the data — the exact thing we wanted to fix — and this only works with small integer weights.
Instead, I found a nice “regular” k-means example from stackexchange (thanks Gareth Rees!) which I stripped down and modified to be observation-weighted:
import randomimport numpy as npimport pandas as pd import scipy.spatialfrom haversine import haversinedef distance(p1,p2):  return haversine(p1[1:],p2[1:])def cluster_centroids(data, clusters, k):  results=[]  for i in range(k):    results.append( np.average(data[clusters == i],weights=np.squeeze(np.asarray(data[clusters == i][:,0:1])),axis=0))  return resultsdef kmeans(data, k=None, centroids=None, steps=20):  # Forgy initialization method: choose k data points randomly.  centroids = data[np.random.choice(np.arange(len(data)), k, False)]  for _ in range(max(steps, 1)):    sqdists = scipy.spatial.distance.cdist(centroids, data, lambda u, v: distance(u,v)**2)    # Index of the closest centroid to each data point.    clusters = np.argmin(sqdists, axis=0)    new_centroids = cluster_centroids(data, clusters, k)    if np.array_equal(new_centroids, centroids):      break    centroids = new_centroids    return clusters, centroids#setupdata = pd.read_csv(“us_census.csv”)data = data[~data[‘state’].isin([‘AK’,’HI’,’PR’])]vals = data[[‘population’,’latitude’,’longitude’]].valuesk = 3random.seed(42)#run itclusters,centroids=observation_weighted_kmeans(vals,k)#outputdata[‘c’]=[int(c) for c in clusters]lats = [centroids[i][1] for i in range(k)]data[‘clat’] = data[‘c’].map(lambda x: lats[x])longs = [centroids[i][2] for i in range(k)]data[‘clong’] = data[‘c’].map(lambda x: longs[x])data.to_csv("clustered_us_census.csv", index=False)
On a old Macbook Air, run time ranged from 2 seconds (k=1) to 160 seconds (k=48).
Where is the center of mass of the U.S. population (i.e., k=1)? That appears to be in south east Missouri.
And for k=10,
Finally, if we were to “redistrict” our states based on population density then we see a very different picture. At least 5 states would be divvied out to neighbors and CA and TX split into four or five regions.
Update (2017–05–09): here is the full source code: https://github.com/leapingllamas/medium_posts/tree/master/observation_weighted_kmeans
Update (2018–06–10): there are some bugs in the medium post code shown above. However, the github code works great. To be clear, make sure that you are running: https://github.com/leapingllamas/medium_posts/blob/master/observation_weighted_kmeans/medium.py
P.S. if you are going to be running clustering on some subset of the data, say for districts within a single state, you will need to make sure that you adjust the centroids initialization. That is, ensure that the initial centroids are randomly distributed across the area of interest (not across the broader US as in the example code) otherwise they will end up in a single cluster.