将循环中的 where 应用到不同的列集的更有效方法 - python

Question

I want cross tabs for columns 'month', 'week' and 'year' against column 'a', but i only want to replace values 0 and 1 for columns 'month' and 'week'. I have a code below that does technically work, but i was wondering if there was a more efficient way of writing it? any pointers would be great! thank you

import pandas as pd
import numpy as np 

out = {}
df = pd.DataFrame({'a': ['a','b','b','a','b','b','a','b','a'], 
                   'month':['march','march','january', 'march','january','january', 'may','march','march'],
                   'week':['1','1','1', '1','2','3', '3','2','1'],
                   'year':['5','3','4', '3','1','1', '1','1','1']})

cols_a =['month', 'week']
cols_b = ['year']

out1 = {}
out2={}

for col in cols_a:
    ct1 = pd.crosstab(df.a, df[col])
    ct2 = pd.DataFrame(ct1.where(ct1 >=2, 'group_a'))
    out1[f'{(col)}'] = ct2

for col in cols_b:
    ct3 =  pd.crosstab(df.a, df[col])
    out2[f'{(col)}'] = ct3

out3 = {**out1, **out2}

the current output looks like this, which is correct

{'month': month  january  march      may
 a                             
 a      group_a      3  group_a
 b            3      2  group_a,
 'week': week  1        2        3
 a                        
 a     3  group_a  group_a
 b     2        2  group_a,
 'year': year  1        3        4        5
 a                                 
 a     2  1  0  1
 b     3  1  1  0}

Answer 1

Idea is join list and use where only for columns if exist in cols_a :

cols_a = ['month', 'week']
cols_b = ['year']

out = {}
for col in cols_a + cols_b:
    ct1 = pd.crosstab(df.a, df[col])
    if col in cols_a:
        ct1 = ct1.where(ct1 >=2, 'group_a')
    out[f'{(col)}'] = ct1

Answer 2

Not revolutionary different, but here is a solution in form of a dict comprehension :

{k: pd.crosstab(df['a'], df[k]).applymap(lambda x: 'group_a' if x<2 else x)
    if k in cols_a else
    pd.crosstab(df['a'], df[k])
 for k in cols_a+cols_b
}

output:

{'month': month  january  march      may
 a                             
 a      group_a      3  group_a
 b            3      2  group_a,
 'week': week  1        2        3
 a                        
 a     3  group_a  group_a
 b     2        2  group_a,
 'year': year  1  3  4  5
 a               
 a     2  1  0  1
 b     3  1  1  0}

Here is an alternative to avoid the applymap :

def group_a(df):
    return df.where(df >= 2, 'group_a')

{k: pd.crosstab(df['a'], df[k]).transform(group_a)
    if k in cols_a else
    pd.crosstab(df['a'], df[k])
 for k in cols_a+cols_b}