Create a dictionary from pandas dataframe

I have the following code

pd.DataFrame(index=["A", "A", "B","B","B"], columns=["col1","col2","col3"],
data=np.array([[5,3,5],[8,6,4],[8,6,4],[2,7,6],[8,6,4]]))

which results in the following dataframe

    col1    col2    col3
A      5     3       5
A      8     6       4
B      8     6       4
B      2     7       6
B      8     6       4

I want to create a dictionary out of it with keys A and B and values to be dataframes with the rest of the data. For example for A the value should be a dataframe like below.

col1    col2    col3
5        3       5
8        6       4

Upvotes: 3

Views: 120

Answers (1)

jezrael
jezrael

Reputation: 862451

Use dictionary comprehension with groupby and DataFrame.reset_index for default index:

d = {k:v.reset_index(drop=True) for k, v in df.groupby(level=0)}
print (d)
{'A':    col1  col2  col3
0     5     3     5
1     8     6     4, 'B':    col1  col2  col3
0     8     6     4
1     2     7     6
2     8     6     4}

print (d['A'])
   col1  col2  col3
0     5     3     5
1     8     6     4

If index values is necessary not changed is possible use:

d = dict(tuple(df.groupby(level=0)))
print (d['A'])
   col1  col2  col3
A     5     3     5
A     8     6     4

But it is same like selecting original data with DataFrame.loc:

print (df.loc['A'])
   col1  col2  col3
A     5     3     5
A     8     6     4

Setup:

df = pd.DataFrame(index=["A", "A", "B","B","B"], 
                  columns=["col1","col2","col3"],
                  data=np.array([[5,3,5],[8,6,4],[8,6,4],[2,7,6],[8,6,4]]))

Upvotes: 2

Related Questions