Cohere Health uses patient feature vectors to support risk classification workflows. Implement K-nearest neighbors classification for a batch of query vectors without relying on machine-learning libraries.
Given labeled training vectors, return one predicted label for each query vector. Use Euclidean distance and select the k closest training points. When weighted is false, each neighbor contributes one vote. When it is true, each neighbor contributes 1 / distance weight. If a query exactly matches one or more training vectors, use only those zero-distance points for voting. Break every voting tie by choosing the lexicographically smallest label.
Implement knn_predict(train_X, train_y, query_X, k, weighted). train_X and query_X are lists of equal-dimensional numeric vectors. train_y contains one non-empty string label for each training vector. Return a list of strings, one per query vector. The input must not be mutated.
Use a bounded max-heap so neighbor selection requires O(n log k) time per query rather than sorting every distance.
def knn_predict(train_X, train_y, query_X, k, weighted):