Implement a k-nearest neighbors algorithm.

In this interview, Angie asks Raj (Snapchat) to "implement the K-nearest neighbor algorithm."

Step 1: Understand the problem

Having asked the appropriate clarifying questions, we’ll make the following assumptions:

• We are implementing a binary classifier.
• We will use the Euclidean distance as the metric to compute the distance between an input data point and data points in the training set.
• The input features, X, will be a NumPy array with arbitrary dimension N examples and D features.

Step 2: Discuss the approach

KNN for classification will involve predicting an incoming data point based on the most prominent label amongst the data points in the training set that are closest to it.

As we discussed previously, we will be using the Euclidean distance to determine “closeness.”

We’ll define two functions, one to calculate the distance and one to infer the KNN algorithm: we can call these “euclidean_distance” and “knn.”

Step 3: Implement the algorithm

Python
def euclidean_distance(x, y):
  return np.sqrt((x - y).T.dot(x - y))

def knn(X_train, y_train, X_new, k):
  distances = []

As clarified before, we will be using the Euclidean distance to determine closeness, and we’ll assume that we have the training input data and labels. In this case, we will also assume that the “knn” function takes in a single data point that represents input features. This could be extended to doing KNN for multiple data points.

We’ll start by implementing the basic logic for KNN, and then implement the logic for Euclidean distance afterwards. The algorithm involves calculating the distances between all the training data points and the incoming data point. For this, we’ll use a list to maintain the distances and populate the list using the distance function and a “for” loop:

Python
distances = []
for X_i in X_train:
   distance = euclidean_distance(X_i, X_new)
   distances.append(distance)

Now, we have to get the indices of the top k minimum distances. We can use the NumPy built-in sorting to do this.

Python
top_k_indices = np.argsort(distances)[:k]
top_k_labels = [y_train[idx].item() for idx in top_k_indices]

The first line will get us the indices of the distances in sorted order. Since NumPy sorts in descending order by default, we will use “[::-1]” to reverse it to ascending order. Then, we will take the top K by indexing this sorted list of indices.

The second line will get the actual label at each of those indices.

Now that we have the labels of the closest data points, we just need to take the majority label and predict it.

Python
# Ensure floating-point division by explicitly converting `k` to a float
  if sum(top_k_labels) > (k / 2.0):
    return 1
  else:
    return 0

Since this is a binary classification, we know whether to predict 1 or 0 based on the sum of the labels. If it exceeds half of the total, the majority label must be 1. Else it is 0.

Now that we’ve implemented the logic for inference, let’s implement our helper function for the Euclidean distance.

The Euclidean distance is calculated by taking the squared difference between each corresponding feature of the two data points, summing them up, and then taking the square root of the entire result. This can actually be written more cleanly in NumPy to vectorize that logic:

Python
return np.sqrt((x1 - x2).T.dot(x1 - x2))

Step 4: Test code and discuss results

We can test our code by creating some dummy data to make sure that our algorithm correctly gets the K nearest neighbors. We’ll choose some training data/labels that we know distances will be closer to than others.

Python
X_train = [
   [0, 0],
   [0, 1],
   [2, 0]
   [2, 2],
   [2, 1]
]

y_train = [0, 0, 0, 1, 1]

Test case 1: X_new = [0, 0], k = 3

The expected result is 0, because the closest points are [0,0] (label of 0), [0, 1] (label of 0), and [2, 0] (label of 0).

Test case 2: X_new = [2,2], k = 3

The expected result is 1, because the closest points are [0, 2] (label of 0), [2,2] (label of 1), and [2,1] (label of 1)