Python Pandas：有效计算值大于或等于一组值的行数，按键列分组

Question

假设我有两个 Pandas 数据帧：

• df1具有列k1 ( str )、 k2 ( str ) 和v ( float )，以及
• df2有一列w ( float )。

我可以假设行df1已排序，首先按k1 ，然后按k2 ，最后按v 。 我可以假设df2中w的值是唯一且已排序的。

我的目标是创建一个新的 DataFrame df3 ，其中包含k1 、 k2 、 w和count_ge 。 DataFrame df3应该为k1 、 k2和w的每个唯一组合有一行； 列count_ge应该是df1中具有相同k1和k2值的行数，并且v的值大于或等于w的值。

以下代码是一个幼稚的实现，似乎可以满足我的要求。 有没有一种有效的方法来执行相同的操作？ 理想情况下，代码还应该推广到df1中的两个以上的键。

import pandas as pd

# Generate some example data.
df1 = pd.DataFrame(
    (
        ('A', 'A', 1),
        ('A', 'A', 1),
        ('A', 'A', 3),
        ('A', 'A', 4),
        ('B', 'C', 2),
        ('B', 'C', 6),
    ),
    columns=('k1', 'k2', 'v'),
)

df2 = pd.DataFrame(
    (0, 2, 5),
    columns=('w',),
)

# Get all unique combinations of k1, k2, and w.
# In Pandas 1.2.0, we can use `merge(how='cross')` for this instead.
df3 = (
    df1[['k1', 'k2']]
    .drop_duplicates()
    .assign(_key=1)
    .merge(df2.assign(_key=1), on='_key')
    .drop(columns='_key')
)

# For each row in df3, count the number of rows in df1 that have the same values of k1 and k2,
# and a value of v that is greater than or equal to w.
df3['count_ge'] = 0
for i, (k1, k2, w, _) in df3.iterrows():
    df3.loc[i, 'count_ge'] = len(df1.query(f'k1 == {k1!r} and k2 == {k2!r} and v >= {w!r}'))
df3

Answer 1

使用交叉merge初始化df3 ：

df3 = df1[["k1", "k2"]].drop_duplicates().merge(df2, how='cross')
>>> df3
  k1 k2  w
0  A  A  0
1  A  A  2
2  A  A  5
3  B  C  0
4  B  C  2
5  B  C  5

然后对于count_ge列，您可以像这样使用lambda function ：

df3['count_ge'] = df3.apply(lambda x: df1[(df1["k1"]==x["k1"])&(df1["k2"]==x["k2"])&(df1["v"]>=x["w"])].shape[0], axis=1)
>>> df3
  k1 k2  w  count_ge
0  A  A  0         4
1  A  A  2         2
2  A  A  5         0
3  B  C  0         2
4  B  C  2         2
5  B  C  5         1

Answer 2

另一种可能的方法是使用np.histogram 。 这种方法看起来相当干净，但存在复制数据帧的潜在缺点pd.concat 。 仍然欢迎其他建议。

import numpy as np
import pandas as pd

# Generate some example data.
df1 = pd.DataFrame(
    (
        ('A', 'A', 1),
        ('A', 'A', 1),
        ('A', 'A', 3),
        ('A', 'A', 4),
        ('B', 'C', 2),
        ('B', 'C', 6),
    ),
    columns=('k1', 'k2', 'v'),
)

df2 = pd.DataFrame(
    (0, 2, 5),
    columns=('w',),
)

# For each unique combination of (k1, k2, w), count the number of rows in df1 that have the same values of k1 and k2,
# and a value of v that is greater than or equal to w.
# such
v_bins = np.concatenate((df2['w'], [np.inf]))
df3s = []
for (k1, k2), v in df1.groupby(['k1', 'k2'])['v']:
    df = df2.copy()
    df['count_ge'] = np.histogram(a=v, bins=v_bins)[0][::-1].cumsum()[::-1]
    df['k1'] = k1
    df['k2'] = k2
    df3s.append(df[['k1', 'k2', 'w', 'count_ge']])
pd.concat(df3s)