Artificial intelligence → Module 3 → The nearest neighbor algorithm (k-NN)

The nearest neighbor algorithm (k-NN) from scratch

How would you guess the type of an unknown fruit? You would probably compare it with fruits you know and say: “it looks most like apples”. The k-nearest neighbors algorithm (k-NN) works on this idea. We will write it ourselves, using only lists and loops.

The idea of the algorithm

1. Each fruit is a point on the chart (x = mass, y = length).
2. For a new fruit we calculate the distance to every known fruit.
3. We take the k nearest (for example k = 3).
4. They “vote”: the type that most of them have is our answer.

We calculate the distance between the points (x₁, y₁) and (x₂, y₂) using the Pythagorean theorem: √((x₁ − x₂)² + (y₁ − y₂)²).

An important step: scaling

The mass is in hundreds of grams (115–250), and the length in centimeters (7–21). A difference of 30 g would “count” much more in the distance than a difference of 10 cm – even though 10 cm is a huge difference for fruit. That is why we convert both features to the range from 0 to 1:

new value = (value − smallest) / (largest − smallest)

This way both features are equally important. You will see what happens without scaling in the first task.

The program

# The k-nearest neighbors algorithm (k-NN) – without libraries, only lists and loops
import csv
import math
import random

# 1) loading the data
data = []
with open("fruit.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        data.append((float(row["mass"]), float(row["length"]), row["kind"]))

# 2) scaling: mass (100–250) and length (7–21) have different ranges,
#    so we convert both features to the range 0..1 so that they are equally important
min_m = min(p[0] for p in data); max_m = max(p[0] for p in data)
min_d = min(p[1] for p in data); max_d = max(p[1] for p in data)


def scale(mass, length):
    return ((mass - min_m) / (max_m - min_m), (length - min_d) / (max_d - min_d))


examples = [(scale(m, d), v) for m, d, v in data]

# 3) splitting into training (80 %) and test (20 %) data
random.seed(1)                       # the same “random” order every time the program runs
random.shuffle(examples)
cut = int(len(examples) * 0.8)
train, test = examples[:cut], examples[cut:]


def distance(a, b):
    """The ordinary (Euclidean) distance between two points."""
    return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)


def knn(point, k):
    """Return the type that the majority of the k nearest training examples have."""
    ordered = sorted(train, key=lambda p: distance(point, p[0]))
    votes = {}
    for _, kind in ordered[:k]:
        votes[kind] = votes.get(kind, 0) + 1
    return max(votes, key=votes.get)


# 4) checking the accuracy for several values of k
for k in [1, 3, 5]:
    correct = sum(1 for point, kind in test if knn(point, k) == kind)
    print(f"k = {k}: {correct} of {len(test)} correct ({correct / len(test) * 100:.0f} %)")

# 5) a new fruit
new_fruit = scale(200, 8.3)
print("A fruit of 200 g and 8.3 cm is:", knn(new_fruit, 3))

Program output:

k = 1: 5 of 6 correct (83 %)
k = 3: 5 of 6 correct (83 %)
k = 5: 5 of 6 correct (83 %)
A fruit of 200 g and 8.3 cm is: orange

• random.seed(1) makes the “random” shuffling the same every time the program runs – that way we can compare the results.
• sorted(..., key=lambda p: distance(point, p[0])) sorts the examples from the nearest to the farthest.
• The dictionary votes counts the votes; max(votes, key=votes.get) returns the type with the most votes.
• On the 6 test examples the model makes one mistake – it is an apple that is very close to the oranges.

Bonus: a decision map

If we ask k-NN about every point of the plane, we get a map: which type would be in which place. The borders between the colors are the borders the model has “learned”.

# Bonus: a “decision map” – for every point of the plane k-NN says which type it would be
import csv
import math
import matplotlib.pyplot as plt

data = []
with open("fruit.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        data.append((float(row["mass"]), float(row["length"]), row["kind"]))

min_m = min(p[0] for p in data); max_m = max(p[0] for p in data)
min_d = min(p[1] for p in data); max_d = max(p[1] for p in data)


def scale(mass, length):
    return ((mass - min_m) / (max_m - min_m), (length - min_d) / (max_d - min_d))


examples = [(scale(m, d), v) for m, d, v in data]


def knn(point, k=3):
    ordered = sorted(examples, key=lambda p: math.dist(point, p[0]))
    votes = {}
    for _, kind in ordered[:k]:
        votes[kind] = votes.get(kind, 0) + 1
    return max(votes, key=votes.get)


colors = {"apple": "red", "banana": "gold", "orange": "darkorange"}
# a grid of points: mass from 100 to 260 g, length from 6 to 22 cm
for mass in range(100, 261, 4):
    for d10 in range(60, 221, 4):
        length = d10 / 10
        plt.scatter(mass, length, c=colors[knn(scale(mass, length))], s=12, alpha=0.25, marker="s")

for kind in colors:          # the real examples on top of the map
    plt.scatter([m for m, d, v in data if v == kind], [d for m, d, v in data if v == kind],
                c=colors[kind], edgecolors="black", label=kind)
plt.xlabel("mass (g)")
plt.ylabel("length (cm)")
plt.title("k-NN (k = 3): which type would be in each place")
plt.legend()
plt.show()

k-NN has no “training” in the real sense – it just remembers all the examples. That is why it is simple, but slow when there are a million examples: for each new fruit it has to calculate a million distances.

Common misconceptions

• “k-NN “understands” fruit” – it only measures distances to known examples.
• “A bigger k is always better” – a too big k also asks distant examples; when k is equal to the number of all examples, the answer is always the type that has the most examples.

Practice tasks

Try it yourself first, and only then open the solution.

1. Remove the scaling (use the mass and length as they are) and check what k-NN says for a small apple of 130 g and 8 cm.
2. Write a program that classifies a list of new fruits at once: (150 g, 7.6 cm), (128 g, 18 cm), (240 g, 9.2 cm), (190 g, 8.3 cm).
3. What would k-NN answer for any fruit if k were equal to the number of all training examples (24)? Why?

# Task: k-NN WITHOUT scaling – mass “outvotes” length
import csv, math, random

data = []
with open("fruit.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        data.append(((float(row["mass"]), float(row["length"])), row["kind"]))

random.seed(1)
random.shuffle(data)
train, test = data[:24], data[24:]


def knn(point, k):
    ordered = sorted(train, key=lambda p: math.dist(point, p[0]))
    votes = {}
    for _, v in ordered[:k]:
        votes[v] = votes.get(v, 0) + 1
    return max(votes, key=votes.get)


correct = sum(1 for t, v in test if knn(t, 3) == v)
print(f"Without scaling: {correct} of {len(test)} correct")
# A fruit of 130 g and 8 cm (a small apple) – what does the model without scaling say?
print("130 g, 8 cm ->", knn((130, 8), 3))

Without scaling: 5 of 6 correct
130 g, 8 cm -> banana

# Task: classify several new fruits at once (k-NN with scaling, k = 3)
import csv, math

data = []
with open("fruit.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        data.append((float(row["mass"]), float(row["length"]), row["kind"]))
min_m = min(p[0] for p in data); max_m = max(p[0] for p in data)
min_d = min(p[1] for p in data); max_d = max(p[1] for p in data)


def scale(m, d):
    return ((m - min_m) / (max_m - min_m), (d - min_d) / (max_d - min_d))


examples = [(scale(m, d), v) for m, d, v in data]


def knn(point, k=3):
    ordered = sorted(examples, key=lambda p: math.dist(point, p[0]))
    votes = {}
    for _, v in ordered[:k]:
        votes[v] = votes.get(v, 0) + 1
    return max(votes, key=votes.get)


new_fruits = [(150, 7.6), (128, 18.0), (240, 9.2), (190, 8.3)]
for m, d in new_fruits:
    print(f"{m} g, {d} cm -> {knn(scale(m, d))}")

150 g, 7.6 cm -> apple
128 g, 18.0 cm -> banana
240 g, 9.2 cm -> orange
190 g, 8.3 cm -> orange

Check your knowledge