clustering
k-means
data science
algorithm
machine learning

Group n points in k clusters of equal size

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Grouping n points into k clusters of equal size is not the same problem as ordinary k-means. Standard k-means minimizes distance to centroids, but it does not enforce balanced cluster sizes, so equal-sized clustering becomes a constrained optimization problem.

Why Plain k-Means Does Not Solve It

Vanilla k-means assigns each point to its nearest centroid. That can easily create uneven clusters, especially when the data distribution is skewed.

If the requirement is that every cluster must contain exactly n / k points, then you are adding a hard capacity constraint. That changes the problem substantially. It is no longer enough to say "assign each point to the closest center."

You also need n to be divisible by k if the size must be exactly equal.

A Simple Balanced Assignment Heuristic

A practical way to understand the problem is:

  1. choose candidate centroids
  2. compute distances from every point to every centroid
  3. assign points while respecting remaining capacity per cluster

The code below shows a small greedy example:

python
1import numpy as np
2
3points = np.array([
4    [0.0, 0.0], [0.2, 0.1], [0.1, -0.1],
5    [5.0, 5.0], [5.2, 5.1], [4.9, 4.8]
6])
7
8centroids = np.array([
9    [0.0, 0.0],
10    [5.0, 5.0]
11])
12
13capacity = len(points) // len(centroids)
14assignments = [-1] * len(points)
15remaining = [capacity] * len(centroids)
16
17distances = []
18for i, point in enumerate(points):
19    for j, centroid in enumerate(centroids):
20        distance = np.linalg.norm(point - centroid)
21        distances.append((distance, i, j))
22
23distances.sort()
24
25for _, point_index, cluster_index in distances:
26    if assignments[point_index] == -1 and remaining[cluster_index] > 0:
27        assignments[point_index] = cluster_index
28        remaining[cluster_index] -= 1
29
30print(assignments)

This is easy to run and reason about, but it is only a heuristic. It does not guarantee the globally best balanced solution.

Better Approaches for Real Problems

For serious equal-size clustering, people often use:

  • min-cost flow or assignment formulations
  • integer programming
  • balanced k-means variants that alternate centroid updates with constrained reassignment

These methods treat cluster capacity as part of the optimization instead of as an afterthought. They are more accurate, but they also cost more computation.

That tradeoff matters. If the dataset is small and the size constraint is strict, an exact solver may be appropriate. If the dataset is large, a heuristic or approximate method may be the only practical option.

Equal Size Versus Good Geometry

Balanced clustering is often driven by business rules rather than by pure geometric similarity. For example, you may want:

  • equal workloads across teams
  • equally sized customer cohorts
  • balanced partitions for distributed processing

Those are valid reasons, but they can pull against the natural structure of the data. If the data clearly forms one large group and several smaller ones, forcing equal cluster sizes may create unnatural assignments.

So the right question is not only "how do I cluster" but also "how important is the equal-size constraint compared with cluster quality."

Common Pitfalls

  • Running plain k-means and hoping the clusters come out balanced by accident.
  • Forgetting that exact equal-size clustering requires n to be divisible by k.
  • Using a greedy heuristic and assuming it is globally optimal.
  • Enforcing equal sizes even when the data distribution clearly argues against it.
  • Treating the problem as a parameter tweak instead of as a constrained optimization task.

Summary

  • Equal-size clustering is different from standard k-means because cluster capacity is a hard constraint.
  • Plain nearest-centroid assignment does not enforce balanced cluster sizes.
  • Greedy heuristics are easy to implement, but they do not guarantee optimality.
  • Higher-quality approaches usually rely on assignment, flow, or integer-programming ideas.
  • Decide explicitly whether equal size is a true requirement or just a preference that may hurt cluster quality.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.