根据前缀将列加载到多个DataFrame中

Question

I want to load columns with specific prefixes into separate DataFrames.

The columns I want have specific prefixes ie

   A_1 A_2 B_1 B_2 C_1 C_2
   1   0   0   0   0   0
   1   0   0   1   1   1
   0   1   1   1   1   0

I have a list of all the prefixes:

prefixes = ["A", "B", "C"]

I want to do something like this:

for prefix in prefixes:
    f"df_{prefix}" = pd.read_csv("my_file.csv",
                                 usecols=[f"{prefix}_1,
                                          f"{prefix}_2,
                                          f"{prefix}_3,])

So each DataFrame has the prefix in the name, but I'm not quite sure of the best way to do this or the syntax required.

Answer 1

You could try it with a different approach. Load the full csv once. Create three dfs out of it by dropping the columns don't mach your prefix.

x = pd.read_csv("my_file.csv")
notA = [c for c in x.columns if 'A' not in c]
notB = [c for c in x.columns if 'B' not in c]
notC = [c for c in x.columns if 'C' not in c]
a = x.drop(notA,1)
b = x.drop(notB,1)
c = x.drop(notC,1)

Answer 2

Considering you have a big dataframe like this:

In [1341]: df
Out[1341]: 
   A_1  A_2  B_1  B_2  C_1  C_2
0    1    0    0    0    0    0
1    1    0    0    1    1    1
2    0    1    1    1    1    0

Have a master list of prefixes:

In [1374]: master_list = ['A','B','C']

Create an empty dictionary to hold multiple subsets of dataframe:

In [1377]: dct = {}

Loop through the master list and store the column names in the above dict:

In [1378]: for i in master_list:
      ...:     dct['{}_list'.format(i)] = [e for e in df.columns if e.startswith('{}'.format(i))]

Now, the dct has below keys with values:

A_list : ['A_1', 'A_2']
B_list : ['B_1', 'B_2']
C_list : ['C_1', 'C_2']

Then, subset your dataframes like below:

In [1381]: for k in dct:
      ...:     dct[k] = df[dct[k]]

Now, the dictionary has actual rows of dataframe against every key:

In [1384]: for k in dct:
      ...:     print dct[k]

In [1347]: df_A
Out[1347]: 
   A_1  A_2
0    1    0
1    1    0
2    0    1

In [1350]: df_B
Out[1350]: 
   B_1  B_2
0    0    0
1    0    1
2    1    1

In [1355]: df_C
Out[1355]: 
   C_1  C_2
0    0    0
1    1    1
2    1    0

Answer 3

First filter out not matched columns with startswith with boolean indexing and loc , because filter columns:

print (df)
   A_1  A_2  B_1  B_2  C_1  D_2
0    1    0    0    0    0    0
1    1    0    0    1    1    1
2    0    1    1    1    1    0

prefixes = ["A", "B", "C"]
df = df.loc[:, df.columns.str.startswith(tuple(prefixes))]
print (df)
   A_1  A_2  B_1  B_2  C_1
0    1    0    0    0    0
1    1    0    0    1    1
2    0    1    1    1    1

Then create Multiindex by split and then dictionary with groupby for dictioanry of DataFrames:

df.columns = df.columns.str.split('_', expand=True)
print (df)

   A     B     C
   1  2  1  2  1
0  1  0  0  0  0
1  1  0  0  1  1
2  0  1  1  1  1

d = {k: v[k] for k, v in df.groupby(level=0, axis=1)}
print (d['A'])
   1  2
0  1  0
1  1  0
2  0  1

Or use lambda function with split :

d = {k: v for k, v in df.groupby(lambda x: x.split('_')[0], axis=1)}
print (d['A'])
   A_1  A_2
0    1    0
1    1    0
2    0    1