Pandas数据帧使用行特定的N查找每行的最大N个元素

时间:2017-05-19 14:42:02

标签: pandas dataframe max rowwise

我有一个DataFrame:

>>> df = pd.DataFrame({'row1' : [1,2,np.nan,4,5], 'row2' : [11,12,13,14,np.nan], 'row3':[22,22,23,24,25]}, index = 'a b c d e'.split()).T
>>> df
         a     b     c     d     e
row1   1.0   2.0   NaN   4.0   5.0
row2  11.0  12.0  13.0  14.0   NaN
row3  22.0  22.0  23.0  24.0  25.0

和一个系列,它指定每行所需的前N个值

>>> n_max = pd.Series([2,3,4])

Panda使用dfn_max查找每个元素的最大N元素的方式是什么(与随机选择断开关系,就像.nlargest()那样)?

所需的输出是

         a     b     c     d     e
row1   NaN   NaN   NaN   4.0   5.0
row2   NaN  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0

我知道如何在所有行中使用统一/固定N(例如,N = 4)。请注意第3行中的打破平局:

>>> df.stack().groupby(level=0).nlargest(4).unstack().reset_index(level=1, drop=True).reindex(columns=df.columns)
         a     b     c     d     e
row1   1.0   2.0   NaN   4.0   5.0
row2  11.0  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0

但目标再次是具有行特定N。循环显示每行显然不计算(出于性能原因)。而且我尝试使用带掩码的.rank()但是在那里打破平局......

2 个答案:

答案 0 :(得分:1)

根据@ ScottBoston对OP的评论,可以使用以下基于排名的掩码来解决这个问题:

>>> n_max.index = df.index
>>> df_rank = df.stack(dropna=False).groupby(level=0).rank(ascending=False, method='first').unstack()
>>> selected = df_rank.le(n_max, axis=0)
>>> df[selected]
         a     b     c     d     e
row1   NaN   NaN   NaN   4.0   5.0
row2   NaN  12.0  13.0  14.0   NaN
row3  22.0   NaN  23.0  24.0  25.0

答案 1 :(得分:0)

为表现,我会建议NumPy -

def mask_variable_largest_per_row(df, n_max):
    a = df.values
    m,n = a.shape

    nan_row_count = np.isnan(a).sum(1)
    n_reset = n-n_max.values-nan_row_count
    n_reset.clip(min=0, max=n-1, out = n_reset)

    sidx = a.argsort(1)
    mask = n_reset[:,None] > np.arange(n)

    c = sidx[mask]
    r = np.repeat(np.arange(m), n_reset)
    a[r,c] = np.nan
    return df

示例运行 -

In [182]: df
Out[182]: 
         a     b     c     d     e
row1   1.0   2.0   NaN   4.0   5.0
row2  11.0  12.0  13.0  14.0   NaN
row3  22.0  22.0   5.0  24.0  25.0

In [183]: n_max = pd.Series([2,3,2])

In [184]: mask_variable_largest_per_row(df, n_max)
Out[184]: 
       a     b     c     d     e
row1 NaN   NaN   NaN   4.0   5.0
row2 NaN  12.0  13.0  14.0   NaN
row3 NaN   NaN   NaN  24.0  25.0

进一步提升:引入numpy.argpartition替换numpy.argsort应该有所帮助,因为我们并不关心要重置为{{的指数的顺序1}}。因此,基于NaNs的将是 -

numpy.argpartition

运行时测试

其他方法 -

def mask_variable_largest_per_row_v2(df, n_max):
    a = df.values
    m,n = a.shape

    nan_row_count = np.isnan(a).sum(1)
    n_reset = n-n_max.values-nan_row_count
    n_reset.clip(min=0, max=n-1, out = n_reset)

    N = (n-n_max.values).max()
    N = np.clip(N, a_min=0, a_max=n-1)

    sidx = a.argpartition(N, axis=1) #sidx = a.argsort(1)
    mask = n_reset[:,None] > np.arange(n)

    c = sidx[mask]
    r = np.repeat(np.arange(m), n_reset)
    a[r,c] = np.nan
    return df

验证和时间安排 -

def pandas_rank_based(df, n_max):
    n_max.index = df.index
    df_rank = df.stack(dropna=False).groupby(level=0).rank\
               (ascending=False, method='first').unstack()
    selected = df_rank.le(n_max, axis=0)
    return df[selected]

内置大熊猫的 In [387]: arr = np.random.rand(1000,1000) ...: arr.ravel()[np.random.choice(arr.size, 10000, replace=0)] = np.nan ...: df1 = pd.DataFrame(arr) ...: df2 = df1.copy() ...: df3 = df1.copy() ...: n_max = pd.Series(np.random.randint(0,1000,(1000))) ...: ...: out1 = pandas_rank_based(df1, n_max) ...: out2 = mask_variable_largest_per_row(df2, n_max) ...: out3 = mask_variable_largest_per_row_v2(df3, n_max) ...: print np.nansum(out1-out2)==0 # Verify ...: print np.nansum(out1-out3)==0 # Verify ...: True True In [388]: arr = np.random.rand(1000,1000) ...: arr.ravel()[np.random.choice(arr.size, 10000, replace=0)] = np.nan ...: df1 = pd.DataFrame(arr) ...: df2 = df1.copy() ...: df3 = df1.copy() ...: n_max = pd.Series(np.random.randint(0,1000,(1000))) ...: In [389]: %timeit pandas_rank_based(df1, n_max) 1 loops, best of 3: 559 ms per loop In [390]: %timeit mask_variable_largest_per_row(df2, n_max) 10 loops, best of 3: 34.1 ms per loop In [391]: %timeit mask_variable_largest_per_row_v2(df3, n_max) 100 loops, best of 3: 5.92 ms per loop 相当不错!