基于已知最大列和列顺序约束的阈值筛选数据帧

Question

我有一个(7985, 3048)形的数据框，格式如下：

image_name    col1_ID    col2_ID    col3_ID    ...  colN_ID      # header
rdr_001_101   0.00342    0.00015    0.10094    ...  0.34210      # prob. distrib
rdr_001_133   0.00432    0.00025    0.01465    ...  0.22108      # prob. distrib
...
...
rdr_003_167   0.07543    0.00576    0.07523    ...  0.73421      # prob. distrib

现在，对于每个图像，我只想从这些3047个值中选择100个列及其相应的索引。应根据列中的值选择100列（即选择具有最大值的列，直到我们达到100的计数）

约束：列的顺序很重要（即输出中的顺序应与输入中的顺序相同）。请参见下面的输入/输出示例

我有一个基于循环的解决方案（使用df.iterrows），但速度非常慢，而且通常不推荐使用这种方法。此外，这是一个简单的示例，而我的实际数据帧可能有更多行

其他可能（更快）的替代方案是什么

---

示例输入，仅包含5列：

image_name    col1_ID    col2_ID    col3_ID   col4_ID   col5_ID
rdr_001_101   0.00342    0.00015    0.10094   0.34210   0.27651
rdr_001_133   0.43211    0.00025    0.01465   0.22108   0.00123
rdr_003_167   0.07543    0.25761    0.07523   0.01148   0.00341

预期输出，以3列作为阈值：

col3_ID    col4_ID    col5_ID
col1_ID    col3_ID    col4_ID
col1_ID    col2_ID    col3_ID

Answer 1

使用^{}和1d数组，按列的长度首先按np.arange(len(df.columns))，然后索引前N个值，按np.sort和最后一个索引列名称cols排序：

df = df.set_index('image_name')
cols = df.columns.to_numpy()
N = 3

print (np.arange(len(df.columns)))
[0 1 2 3 4]

print (np.arange(len(df.columns))[np.argsort(-df.to_numpy())[:, :N]])
[[3 4 2]
 [0 3 2]
 [1 0 2]]


arr = np.sort(np.arange(len(df.columns))[np.argsort(-df.to_numpy())[:, :N]], axis=1)
print (arr)
[[2 3 4]
 [0 2 3]
 [0 1 2]]

c = cols[arr]
print (c)
[['col3_ID' 'col4_ID' 'col5_ID']
 ['col1_ID' 'col3_ID' 'col4_ID']
 ['col1_ID' 'col2_ID' 'col3_ID']]

如有必要，最后将输出转换为DataFrame：

c1 = [f'top{x+1}' for x in np.arange(N)]
df1 = pd.DataFrame(c, index=df.index, columns=c1)
print (df1)
                top1     top2     top3
image_name                            
rdr_001_101  col3_ID  col4_ID  col5_ID
rdr_001_133  col1_ID  col3_ID  col4_ID
rdr_003_167  col1_ID  col2_ID  col3_ID

---

如果可能，不需要对列名称进行排序1d array：

df = df.set_index('image_name')
cols = df.columns.to_numpy()
c = np.sort(cols[np.argsort(-df.to_numpy())[:, :N]], axis=1)
print (c)

[['col3_ID' 'col4_ID' 'col5_ID']
 ['col1_ID' 'col3_ID' 'col4_ID']
 ['col1_ID' 'col2_ID' 'col3_ID']]