python - Pandas 每行获得前 n 列

Question

我有以下数据框:

df1 = pd.DataFrame(data={'1': ['a', 'd', 'g', 'j'], 
                         '2': ['b', 'e', 'h', 'k'], 
                         '3': ['c', 'f', 'i', 'l'],
                         'top_n': [1, 3, 2, 1]},
                   index=pd.Series(['ind1', 'ind2', 'ind3', 'ind4'], name='index'))
>>> df1
       1  2  3   top_n
index
ind1   a  b  c    1
ind2   d  e  f    3
ind3   g  h  i    2
ind4   j  k  l    1

如何根据 top_n 列只获取每一行的前 N ​​个值？

>>> df1
       1  2    3      top_n
index
ind1   a  NaN  NaN     1
ind2   d  e    f       3
ind3   g  h    NaN     2
ind4   j  NaN  NaN     1

在此示例中，ind3 具有 g 和 h，因为 top_n 值为 2。

Answer 1

使用 numpy对于 bool 掩码:

import numpy as np

a = np.ones_like(df1).cumsum(1)

mask = ((a <= df1['top_n'].values[:,None]) | (a == df1.shape[1]))

out = df1.where(mask)

输出:

       1    2    3  top_n
index                    
ind1   a  NaN  NaN      1
ind2   d    e    f      3
ind3   g    h  NaN      2
ind4   j  NaN  NaN      1

掩码:

array([[ True, False, False,  True],
       [ True,  True,  True,  True],
       [ True,  True, False,  True],
       [ True, False, False,  True]])