Subset column names with f-string and loop-CodePudding

I'm trying to subset column names into a range of lists and I can do it one by one but given I want to do more than 10, I am trying to use a loop with f-string. But mine didn't seem to work. How can I do this with f-string and a loop?

a1_cols = [col for col in df.columns if 'a1_' in col]
a2_cols = [col for col in df.columns if 'a2_' in col]
a3_cols = [col for col in df.columns if 'a3_' in col]
...
a10_cols = [col for col in df.columns if 'a10_' in col]


for i in range(1, 12):
    f'a{i}_cols' = [col for col in df.columns if f'a{i}_' in col]

CodePudding user response：

Create dictionary:

d = {}
for i in range(1, 12):
     d[f'a{i}_cols'] = [col for col in df.columns if f'a{i}_' in col]

Or:

d = {f'a{i}_cols': [col for col in df.columns if f'a{i}_' in col] for i in range(1, 12)}

And then for selecting use:

print (d['a10_cols'])

Possible solution, but not recommended with globals:

for i in range(1, 12):
    globals()[f'a{i}_cols'] = [col for col in df.columns if f'a{i}_' in col]

CodePudding user response：

Just for an alternative way by using regular expression.

import re

col = ["a1_", "a2_", "a3", "a4_", "a5", "b1_", "b2_", 'b3_', "ax_", "ay_"]
find_col = re.compile(r"^a[0-9] \_")

> [x for x in col  if find_col.search(x) != None]

['a1_', 'a2_', 'a4_']