combining groups of columns with boolean values to create multiple new columns

Viewed 119

I was wondering how I can create new columns based on old ones by multiplying two groups of them in my dataframe. So if I have [A, B] and [X, Y]:

   A     B      X     Y
0  True  False  True  False

Now I want this to result in 4 extra columns, of which only the one that has both true is true:

   A     B      X     Y      AX     BX      AY     BY
0  True  False  True  False  True   False   False  False   

In my situation, the groups are bigger, so I'm looking for a solution that can be used by using the two lists of column headers [A, B] and [X, Y]. I have tried list comprehension but I cannot get it to work :-(

1 Answers

Use itertools.product to get the cartesian product of column names then use Series.mul inside a list comprehension to create corresponding column products, finally use pd.concat to concat these products with df:

from itertools import product

l1, l2 = ['A', 'B'], ['X', 'Y']

c = [df[a].mul(df[b]).rename(''.join([a, b])) for a, b in product(l1, l2)]
df = pd.concat([df] + c, axis=1)

Result:

      A      B     X      Y    AX     AY     BX     BY
0  True  False  True  False  True  False  False  False
Related