K-Nearest Neighbors Search [medium]

v1 pairwise distances 로 (N, M) 거리 행렬을 얻었다면, 보통 다음 단계는 각 query 에 대해 k 최근접 을 뽑는 것. 전체 argsort 는 $O(M \log M)$ per row; argpartition 은 $O(M)$ per row — 정확도는 같지만 더 빠름.

알고리즘

1. D = pairwise Euclidean distances shape (N, M).
2. 각 row 에서 가장 작은 k개의 인덱스 추출: idx_unsorted = np.argpartition(D, k, axis=1)[:, :k].
3. 이 인덱스를 거리 순으로 정렬 하고 싶으면, 해당 거리를 가져와 argsort 로 재정렬.

과제

함수 knn_search(X, Y, k) 를 완성하세요.

• X shape (N, d) query, Y shape (M, d) database, k 정수.
• 반환: (idx, dists) 튜플.
  • idx shape (N, k): 각 row 는 query i 의 k 개 nearest Y 인덱스 (거리 오름차순).
  • dists shape (N, k): 해당 거리 (오름차순).
• 루프 금지.

힌트

D = pairwise_euclidean(X, Y)                 # (N, M)
idx_unsorted = np.argpartition(D, k, axis=1)[:, :k]
dist_unsorted = np.take_along_axis(D, idx_unsorted, axis=1)
order = np.argsort(dist_unsorted, axis=1)    # 거리 기준 재정렬
idx = np.take_along_axis(idx_unsorted, order, axis=1)
dists = np.take_along_axis(dist_unsorted, order, axis=1)
return idx, dists

테스트 케이스