Java: k-nearest-neighbors classification

Before you start

Java arrays, sorting with a comparator, and maps.

Classify a point by letting its k closest training points vote.
Predict the sample output and explain the state changes.
Complete the practice task and check the boundary cases.

The concept

k-nearest neighbours is a simple classifier with no training phase. To classify a new point it finds the k training points closest to it, usually by Euclidean distance, and predicts the most common label among them.

The program sorts the indexes of the training points by their distance to the query and counts the labels of the first k. Small k can follow noise, and large k blurs the boundary between classes. Features measured on very different scales should be normalised first, or the large-scale feature will dominate the distance.

Code example

Java

import java.util.*;

class Main {
  static double[][] points = {{1, 1}, {1, 2}, {2, 1}, {6, 6}, {7, 6}, {6, 7}};
  static String[] labels = {"A", "A", "A", "B", "B", "B"};

  static double dist(double[] a, double[] b) {
    return Math.hypot(a[0] - b[0], a[1] - b[1]);
  }

  static Map<String, Integer> classify(double[] query, int k) {
    Integer[] order = {0, 1, 2, 3, 4, 5};
    Arrays.sort(order, Comparator.comparingDouble(i -> dist(points[i], query)));
    Map<String, Integer> votes = new TreeMap<>();
    for (int i = 0; i < k; i++) {
      votes.merge(labels[order[i]], 1, Integer::sum);
    }
    return votes;
  }

  public static void main(String[] args) {
    System.out.println("near (5,5): " + classify(new double[] {5, 5}, 3));
    System.out.println("near (2,2): " + classify(new double[] {2, 2}, 3));
    System.out.println("near (3,3) with k=5: " + classify(new double[] {3, 3}, 5));
  }
}

Output

near (5,5): {B=3}
near (2,2): {A=3}
near (3,3) with k=5: {A=3, B=2}

Run code Example details

Follow the code

Open the walkthrough to inspect the example one step at a time.

Show guided walkthrough

Guided walkthroughReady

1import java.util.*;

3class Main {

4 static double[][] points = {{1, 1}, {1, 2}, {2, 1}, {6, 6}, {7, 6}, {6, 7}};

5 static String[] labels = {"A", "A", "A", "B", "B", "B"};

7 static double dist(double[] a, double[] b) {

8 return Math.hypot(a[0] - b[0], a[1] - b[1]);

9 }

11 static Map<String, Integer> classify(double[] query, int k) {

12 Integer[] order = {0, 1, 2, 3, 4, 5};

13 Arrays.sort(order, Comparator.comparingDouble(i -> dist(points[i], query)));

14 Map<String, Integer> votes = new TreeMap<>();

15 for (int i = 0; i < k; i++) {

16 votes.merge(labels[order[i]], 1, Integer::sum);

17 }

18 return votes;

19 }

21 public static void main(String[] args) {

22 System.out.println("near (5,5): " + classify(new double[] {5, 5}, 3));

23 System.out.println("near (2,2): " + classify(new double[] {2, 2}, 3));

24 System.out.println("near (3,3) with k=5: " + classify(new double[] {3, 3}, 5));

25 }

26}

Follow the example

Start to inspect the source, values, and output. These prepared steps explain the published sample; grouped loop steps describe multiple iterations.

View the complete dry-run table

Prepared trace of the published program. Some rows group repeated iterations.
What happens	State after step	Output so far
Euclidean distance between two points.	`—`	—
Sort the training indexes from nearest to farthest.	`order starts with the closest points`	—
Count one vote for the label of each of the k nearest points.	`votes grow for the nearest labels`	—
With k = 5 the vote includes both clusters, because only three points are labelled A, but A still wins 3 to 2.	`votes grow for the nearest labels`	near (5,5): {B=3} near (2,2): {A=3} near (3,3) with k=5: {A=3, B=2}

Behind the scenes

How Java processes this code

The Java compiler produces bytecode. The JVM loads the class and invokes its main method.

Cost and common mistakes

O(n log n) per query for n training points with this full sort. A selection of the k smallest takes O(n).

An even k can produce ties between two classes. Choose an odd k for two classes, or define a tie-break.
Prediction time grows with the training set because every distance is computed.
Do not evaluate on the same points you trained with; the nearest neighbour of a training point is itself.

Try it yourself

Which class wins for (6,5) with k = 3?

Reveal answer and reasoning

B wins: the nearest points are (6,6), (7,6), and (6,7), all labelled B.