import torch


def reserve_nn_scores(similarities, nn_num):
    """Reserver top-k nearest neighbors' similarity scores
    
    Args:
        similarities (Matrix): test_num * train_num
    
    Returns:
        nn_scores (Matrix): only keep the scores of the top-k nearest neighbors
    """
    scores = torch.FloatTensor(similarities)
    # sorted each row, and get the index
    # ascending order
    sorted_scores, sorted_index = torch.sort(scores, 1, descending=True)
    nn_scores = sorted_scores[:, 0:nn_num]
    nn_index = sorted_index[:, 0:nn_num]
    # only accept float32
    nn_scores = torch.zeros(scores.size()).scatter_(1, nn_index, nn_scores)
    # convert float32 to float64
    return nn_scores.numpy().astype(float)