AI011. KNN 分类器

# 整数平方距离(不开方,精确比较),近邻按 (距离,编号) 排序取前 k;
# 表决键=(−票数,类别),两级并列都取小;C++/Java 距离用 64 位。
import sys


def solve() -> None:
    data = sys.stdin.buffer.read().split()
    if not data:
        return
    it = iter(data)
    n, d, k, m = (int(next(it)) for _ in range(4))
    x, y = [], []
    for _ in range(n):
        x.append([int(next(it)) for _ in range(d)])
        y.append(int(next(it)))
    queries = [[int(next(it)) for _ in range(d)] for _ in range(m)]

    out = []
    for q in queries:
        order = []
        for idx, row in enumerate(x):
            dist = 0
            for j in range(d):
                diff = row[j] - q[j]
                dist += diff * diff
            order.append((dist, idx))
        order.sort()
        votes = {}
        for dist, idx in order[:k]:
            votes[y[idx]] = votes.get(y[idx], 0) + 1
        best_label = -1
        best_count = -1
        for label, count in votes.items():
            if count > best_count or (count == best_count and label < best_label):
                best_label = label
                best_count = count
        out.append(str(best_label))
    print("\n".join(out))


if __name__ == "__main__":
    solve()