Top-k Cosine Retrieval [medium]

v1 cosine similarity matrix 는 모든 쌍 의 (N, M) 행렬. 실전 검색 (embedding search, RAG) 에서는 각 query 에 대해 가장 유사한 k개 만 필요:

1. query 와 database 를 각각 L2-normalize → 단위 벡터.
2. S = X̂ @ Ŷ.T 로 (N, M) 유사도.
3. 각 행 top-k (내림차순) 인덱스 + 유사도 반환.

효율 트릭

• 유사도 오름차순 vs 내림차순: argpartition(-S, k) 로 top-k 뽑기 ( -S 에선 최소 = 원래 최대).
• 정렬은 top-k 안에서만.

과제

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

• X (N, d), Y (M, d), 정수 k.
• 반환: (idx, sims) 튜플.
  • idx shape (N, k): 각 query 의 top-k 유사한 Y 인덱스 (유사도 내림차순).
  • sims shape (N, k): 해당 유사도.
• 루프 금지 (num_classes 루프도 없음).
• 영벡터 방어: norm clip 1e-12.

힌트

Xn = X / np.clip(np.linalg.norm(X, axis=1, keepdims=True), 1e-12, None)
Yn = Y / np.clip(np.linalg.norm(Y, axis=1, keepdims=True), 1e-12, None)
S = Xn @ Yn.T                                     # (N, M)
idx_un = np.argsort(-S, axis=1)[:, :k]            # 내림차순 top-k
sims = np.take_along_axis(S, idx_un, axis=1)
return idx_un, sims

테스트 케이스

#	이름	검증
1	shape (N, k) × 2
2	유사도 내림차순
3	k=1 → 최대 유사도 1개
4	X=Y 일 때 자기 자신 top-1 (sim=1)
5	유사도 범위 [-1, 1]
6	영벡터 → NaN 없음