在Pandas中按条件随机抽样DataFrame的行

Question

在Python的Pandas库中，我有一个包含很多行的表格（DataFrame）：

A = pd.DataFrame( { 'key1': [1 , 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 3, ....], 
                    'col1': [ .... ],
                    'col2': [ .... ],
                        ....
                    'col_n': [ .... ],
                    'val1': [0 , 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, ....]} )

每个“键”（key）只有一行或零行的情况，其中“val1”值为1，其他的行“val1”值都是0。

我想创建一个新的表格，里面包含所有“val1==1”的行，以及对应的每个“val1==1”键下随机选取的一行“val1==0”的行。对于那些所有行“val1==0”的键，我就不考虑了。

B = A[ A['val1']== 1] 

这个方法能解决第一部分的问题。但我不太确定怎么实现第二部分。

举个例子：

key1 col1 col2 val1
 1   a     b    0
 1   c     d    0
 1   e     f    1
 1   g     h    0
 2   a     b    0
 2   c     d    1
 2   e     f    0
 3   a     b    0
 3   c     d    0
 3   e     f    0

我想得到：

key1 col1 col2 val1
 1   a     b    0 #randomly subsampled row where val1==0, for key1==1
 1   e     f    1 #row with val1==1 for key1==1
 2   e     f    0 #randomly subsampled row where val1==0, for key1==2
 2   c     d    1 #row with val1==1 for key1==2
                  #no rows for key1==3, because no rows where key1==3 & val1==1

Answer 1

import io
import pandas as pd
import numpy as np

text = """key1 col1 col2 val1
1   a     b    0
1   c     d    0
1   e     f    1
1   g     h    0
2   a     b    0
2   c     d    1
2   e     f    0
3   a     b    0
3   c     d    0
3   e     f    0"""

df = pd.read_csv(io.BytesIO(text), delim_whitespace=True)

def sample(df):
    mask = df.val1 == 0
    if np.all(mask):
        return None
    else:
        idx1 = mask.idxmin()
        idx0 = np.random.choice(mask[mask].index)
        return df.loc[[idx0, idx1]]

df.groupby("key1").apply(sample)

输出结果：

        key1 col1 col2  val1
key1                        
1    0     1    a    b     0
     2     1    e    f     1
2    4     2    a    b     0
     5     2    c    d     1