带总和的itertools product()函数

时间:2019-02-28 13:56:56

标签: python generator itertools

我有一个包含几列的数据框A,我想将所有这些列“与它们自己”求和,以获得数据框B。

A = [col1 col2 col3 
       0    1    2
       1    1    0
      -1    0    1]

和B应该看起来像:

B = [col1+col2  col1+col3  col2+col3
         1          2           3
         2          1           1
        -1          0           1]

基本上,此操作背后的原理完全是itertools.product()函数中嵌入的内容,该函数生成笛卡尔积。

itertools.product('ABCD','xy')-> Ax Ay Bx由Cx Cy Dx Dy

我只需要应用相同的原理并得到:
function_smg('ABCD','xy')-> A + x A + y B + x B + y C + x C + y D + x D + y

我的数据帧很大,所以我负担不起循环,需要迭代器或生成器。 如果没有函数可以解决问题,我该如何构建一个生成器来做到这一点?

非常感谢

2 个答案:

答案 0 :(得分:1)

实际上,有比itertools产品更精确的解决此问题的方法。尝试itertools combinations

import pandas as pd
from itertools import combinations
A = pd.DataFrame({"col1": [0, 1, -1],
              "col2": [1, 1, 0],
              "col3": [2, 0, 1]})

B = pd.DataFrame() #Create an empty dataframe first
for col1, col2 in combinations(A.columns, 2):
    B[f"{col1}+{col2}"] = A[col1] + A[col2] #Create columns one by one.
    #B["{}+{}".format(col1, col2)] = A[col1] + A[col2] (Before python 3.6)

print(B)
#Output:
   col1+col2  col1+col3  col2+col3
0          1          2          3
1          2          1          1
2         -1          0          1

答案 1 :(得分:1)

这是一种方法。您可以使用itertools.combinations从现有列中获取所有长度2的组合:

from itertools import combinations
c = combinations(df.T.values.tolist(), 2)
# [([0, 1, -1], [1, 1, 0]), ([0, 1, -1], [2, 0, 1]), ([1, 1, 0], [2, 0, 1])]

然后将每个元组中的值添加在一起,并加上:

from itertools import starmap
from operator import add

l = [list(starmap(add,zip(i,j))) for i,j in c]
pd.DataFrame(l, index=df.columns).T

    col1  col2  col3
0     1     2     3
1     2     1     1
2    -1     0     1

或者如果numpy也可以选择:

import numpy as np
c = list(combinations(df.T.values.tolist(), 2))
pd.DataFrame(np.array(c).sum(1), index=df.columns).T

    col1  col2  col3
0     1     2     3
1     2     1     1
2    -1     0     1