Problem 380220 · easy · Level 03 Linear Management & Searching

How Many Cones This Hour?

k-nearest neighbours · regression · averaging · tie-breaking

An ice-cream kiosk wants to know how many cones to prepare. Its log of past hours has X_train[i] = [temperature, hour] and the number of cones sold, y_train[i].

Write knn_average(X_train, y_train, queries, k) that returns a list with one prediction (a float) per row of queries: the mean of y_train over the query's k neighbours. The neighbours are the first k training rows when they are ordered by Euclidean distance to the query, nearest first, with equal distances ordered by index (smaller first). If k is larger than the log, every row is a neighbour.

The tests use kiosk_days(n, seed), which returns (X, y) with whole-number rows [temperature in °C, hour of the day] and whole-number cone counts. It is available in your code for experiments.

Examples

Input:  X_train = [[20, 12], [22, 12], [25, 15], [30, 15]], y_train = [40, 50, 80, 120]
        queries = [[21, 12], [28, 15], [26, 14]], k = 2
Output: [45.0, 100.0, 100.0]
Explanation: [21, 12] is 1 away from both of the first two rows: (40 + 50) / 2 = 45.
[28, 15] is nearest to [30, 15] and [25, 15]: (120 + 80) / 2 = 100.

Input:  the same log, queries = [[26, 14], [21, 12]], k = 3
Output: [83.33333333333333, 56.666666666666664]
Explanation: for [26, 14] the squared distances are 40, 20, 2 and 17, so the neighbours are
rows 2, 3 and 1: (80 + 120 + 50) / 3. For [21, 12], rows 0 and 1 (both at 1) and row 2.

Constraints

  • 1 <= len(X_train) <= 500, 0 <= len(queries) <= 300, 1 <= k <= 1000
  • all rows have the same length (1 to 5) and hold whole numbers; floats are compared with a tolerance of 1e-6

Goals

  • Predict a number by averaging the targets of the k most similar training examples
  • Reuse a precise neighbour order with an index tie-break
  • Handle k larger than the training set
Starting Python…