def leave_top_n(M, n):
    # get top n sorted indexes for each column (cols -> sort indexes -> revert -> take top n rows):
    # we're interested in non-0 correlations as 0s are for users w/o data to compute the correlation

    # FIXME ugly cycle, consider replace with more numpythonic idiom
    # iterate through columns / items
    for i in range(M.shape[1]):
        # for each column, attach original row indexes to data
        nonzero_elements = zip(M[:,i].data, __csr_row_indices(M[:,i].indptr))
        # then sort the data and get indexes of everything outside of top-n .data elements
        # reverse sort by 0th -> split to columns -> take 1th -> take after nth
        tail_indices = zip(*sorted(nonzero_elements, key = lambda a : a[0], reverse = True))[1][n:]
        # set those to zero
        M[tail_indices,i] = 0

    # eliminate zeros from the sparse matrix
    M.eliminate_zeros()
For sparse CSR matrix M with float values, mask everything but up to n top elements in each column
dataset/ETHPy150Open ksavenkov/recsys-001/recsys/score.py/leave_top_n