Convert dataframe to array of arrays - pandas

Viewed 526

I'm angling to output an array of arrays from a pandas df. With the df beneath, I want to subset each unique Group to arrays. I'd also hope to produce a separate array for each unique value in id.

import pandas as pd

df = pd.DataFrame({'Int_1': [1.0, 2.0, 1.0, 3.0, 1.0, 2.0, 3.0, 2.0], 
               'Int_2': [1.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 2.0],
               'Period': [1, 1, 1, 1, 2, 2, 2, 2],
               'Group': ['A', 'B', 'A', 'B', 'A', 'B', 'A', 'B'],
               'id': ['1', '2', '3', '4', '1', '2', '3', '4']})

Group_A = [df[df['Group'] == 'A'][['Int_1','Int_2']].to_numpy()]
Group_B = [df[df['Group'] == 'B'][['Int_1','Int_2']].to_numpy()]

print(Group_A)

intended output:

[array([[1.0, 1.0],
       [1.0, 1.0]]), array([[1.0, 2.0], 
       [3.0, 1.0]])]
1 Answers

If need separate array for each Group first filter by boolean indexing and then in list comprehension convert columns to 2d array:

arrA = [g[['Int_1','Int_2']].to_numpy() for i, g in df[df['Group'] == 'A'].groupby('id')]
print (arrA)
[array([[1., 1.],
       [1., 1.]]), array([[1., 2.],
       [3., 1.]])]

If need arrays per Group and id use:

arr = [g[['Int_1','Int_2']].to_numpy() for i, g in df.groupby(['Group', 'id'])]
print (arr)
[array([[1., 1.],
       [1., 1.]]), array([[1., 2.],
       [3., 1.]]), array([[2., 2.],
       [2., 1.]]), array([[3., 2.],
       [2., 2.]])]
Related