带有 Python 和 Pandas 的状态计数机

Question

我正在练习使用 Pandas、Lambda 函数，并面临一项艰巨的任务。我已经有了一个“正式”正确的解决方案，但绝对效率低下。

这就是问题所在：

我有一个 Pandas DataFrame df，它是这样的（在本文末尾生成此示例的代码）：

     id  type
0  1003     G
1  1003     A
2  1002     T
3  1002     A
4  1001     A
5  1003     A
6  1002     G
7  1003     A
8  1001     T
9  1001     A

预期输出：每种不同类型（A、C、G、T）的新列包含唯一 id 的数量，在它们出现在表中的最后一行中，具有该类型。

可能的输出是这样的（经过编辑以匹配所需的结果）：

     id  num_A  num_C  num_G  num_T type
0  1003      0      0      1      0    G
1  1003      1      0      0      0    A
2  1002      1      0      0      1    T
3  1002      2      0      0      0    A
4  1001      3      0      0      0    A
5  1003      3      0      0      0    A
6  1002      2      0      1      0    G
7  1003      2      0      1      0    A
8  1001      1      0      1      1    T
9  1001      2      0      1      0    A

为了达到这个目标，我做了以下事情（如果您将输出与上面显示的表格进行比较，这将无法正常工作）：

1. 定义了一个临时的 DataFrame tmp，它存储了所有可能的 id 的状态（在这个例子中最多有 9 个）：

    id  type_A  type_C  type_G  type_T
   
   
   0  1001       0       0       0       0
   1  1002       0       0       0       0
   2  1003       0       0       0       0
   3  1004       0       0       0       0
   4  1005       0       0       0       0
   5  1006       0       0       0       0
   6  1007       0       0       0       0
   7  1008       0       0       0       0
   8  1009       0       0       0       0
   

2. 定义了一个迭代循环，它检查df 中每一行的类型，然后相应地更新tmp DataFrame 的状态：

代码如下：

for df_row in range(0, df.shape[0]):
    if df.type[df_row] == 'A':
        for tmp_row in range(0, tmp.shape[0]):
            if tmp.id[tmp_row] == df.id[df_row]:
                tmp.type_A[tmp_row] = 1
        df.num_A[df_row] = tmp.type_A.sum()
    if df.type[df_row] == 'C':
        for tmp_row in range(0, tmp.shape[0]):
            if tmp.id[tmp_row] == df.id[df_row]:
                tmp.type_C[tmp_row] = 1
        df.num_C[df_row] = tmp.type_C.sum()
    if df.type[df_row] == 'G':
        for tmp_row in range(0, tmp.shape[0]):
            if tmp.id[tmp_row] == df.id[df_row]:
                tmp.type_G[tmp_row] = 1
        df.num_G[df_row] = tmp.type_G.sum()
    if df.type[df_row] == 'T':
        for tmp_row in range(0, tmp.shape[0]):
            if tmp.id[tmp_row] == df.id[df_row]:
                tmp.type_T[tmp_row] = 1
        df.num_T[df_row] = tmp.type_T.sum()

我想了解的是，如果使用 Lambda 函数或其他方法是否可以获得更快的结果，这在性能方面也更好。

要生成像我这样的示例 DataFrame，您可以使用以下代码（也欢迎提出有关如何调整的建议，以便我了解更多信息）：

df = pd.DataFrame({'id': np.random.randint(1001, 1004, size=10), \
                   'type_tmp': np.random.randint(1, 4, size=10), \
                   'type': '', \
                   'num_G': 0, 'num_A': 0, 'num_T': 0, 'num_C': 0})
for r in range(0, df.shape[0]):
    if df.type_tmp[r] == 1:
        df.type[r] = 'G'
    if df.type_tmp[r] == 2:
        df.type[r] = 'A'
    if df.type_tmp[r] == 3:
        df.type[r] = 'T'
    if df.type_tmp[r] == 4:
        df.type[r] = 'C'
df = df.drop(columns='type_tmp')

临时DataFrame定义如下：

tmp = pd.DataFrame({'id': np.arange(1001, 1010), 'type_A': 0, 'type_C': 0, 'type_G': 0, 'type_T': 0})

感谢您宝贵的时间。

Answer 1

到目前为止，要从数据帧中找到唯一的类型集，您可以从开始到每一行获取数据帧的切片，然后将其强制转换为 Set 并获取长度。如果您已经使用正确的列（全 0）设置了数据框，则可以将这个集合的长度插入到正确的位置：

for index, row in df.iterrows():
    l = len(set(df['type'].head(index)))
    t = row['type']
    df['num_'.format(t)][index] = t

如果有帮助，请告诉我，如果您需要，我可以添加更多。

Answer 2

解决问题的修订版，这与原来的问题有很大不同，我们可以简单地旋转和向前填充以获取任意行的状态，然后使用 value_counts 获取数字：

state = df.reset_index().pivot(index="index", columns="id").ffill()
counts = state.apply(pd.value_counts, axis=1).reindex(["A", "C", "G", "T"], axis=1)
counts = counts.fillna(0).astype(int)
out = df.join(counts)

这给了我

In [193]: out
Out[193]: 
     id type  A  C  G  T
0  1003    G  0  0  1  0
1  1003    A  1  0  0  0
2  1002    T  1  0  0  1
3  1002    A  2  0  0  0
4  1001    A  3  0  0  0
5  1003    A  3  0  0  0
6  1002    G  2  0  1  0
7  1003    A  2  0  1  0
8  1001    T  1  0  1  1
9  1001    A  2  0  1  0