將循環中的 where 應用到不同的列集的更有效方法 - python

Question

我想要“月”、“周”和“年”列與“a”列的交叉表，但我只想替換“月”和“周”列的值 0 和 1。 我下面有一個在技術上可行的代碼，但我想知道是否有更有效的編寫方式？ 任何指針都會很棒！ 謝謝

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}

當前輸出看起來像這樣，這是正確的

{'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

想法是連接列表，如果cols_a存在列，則僅使用where ：

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

沒有革命性的不同，但這里有一個字典理解形式的解決方案：

{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
}

輸出：

{'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}

這是避免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}