python - 如何为每个term保留k个最相似的term记录，并将不太相似的用0替换

Question

我有一个像这样的 pandas 数据框。

它是一个逐个术语的相似度矩阵。对于一个nxn个逐个术语的相似度矩阵，我想保留k个记录，这是每个术语最相似的术语记录，并将其余的n-k个不太相似的记录替换为0。对于这个玩具示例，我想获得一个像这样的数据框

如果 k=2。

frame=pd.DataFrame(data=np.array([[1,0.5,3,0.3],[0.5,1,0.3,0.4],[3,0.3,1,0.7],[0.3,0.4,0.7,1]]),columns=['w1','w2','w3','w4'])
frame.index=['w1','w2','w3','w4']

您能告诉我在将其应用于大矩阵时如何使用 pandas 编写此功能吗？

Answer 1

这是一种利用 NumPy's advanced indexing 的方法和 np.argpartition用于选择和重置每行适当数量的元素 -

k = 2 # no. of records to keep
a = frame.values # Extract the values as an array view
n = a.shape[1] - k # no. of elements to be reset per row
idx = np.argpartition(a,n,axis=1)[:,:n] # smallest n column indices per row
a[np.arange(idx.shape[0])[:,None], idx] = 0 # reset those in array/dataframe

示例运行 -

In [478]: frame=pd.DataFrame(data=np.array([[1,0.5,0.2,0.3],[0.5,1,0.3,0.4],\
     ...:   [0.2,0.3,1,0.7],[0.3,0.4,0.7,1]]),columns=['w1','w2','w3','w4'])
     ...: frame.index=['w1','w2','w3','w4']
     ...: 

In [479]: frame
Out[479]: 
     w1   w2   w3   w4
w1  1.0  0.5  0.2  0.3
w2  0.5  1.0  0.3  0.4
w3  0.2  0.3  1.0  0.7
w4  0.3  0.4  0.7  1.0

## After code run with k=2

In [481]: frame
Out[481]: 
     w1   w2   w3   w4
w1  1.0  0.5  0.0  0.0
w2  0.5  1.0  0.0  0.0
w3  0.0  0.0  1.0  0.7
w4  0.0  0.0  0.7  1.0