根据具有已知最大列和列排序约束的阈值过滤数据帧

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

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

Constraint：列的顺序很重要（即输出中的顺序应与输入中的顺序相同）。请参阅下面的示例 input/output。

我有一个基于循环的解决方案（使用 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

将 numpy.argsort 与 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']]

根据具有已知最大列和列排序约束的阈值过滤数据帧

filtering dataframe based on threshold value with known maximum columns and column ordering constraints

python

performance

memory-efficient

dataframe

pandas