Q:

Pandas how can I locate group n in a sorted groupby?

I have a massive Pandas dataframe converted to a tensor for deep learning. Now I want to access group n in both my dataframe and my array.
E.g., how do I access the dataframe group that eventually becomes group 3 in the array, with a given set of sorting keys?
On a large dataset, converting the dataframe to an array is very slow, as there's roughly 200k groups in millions of rows. So this can only be done once, and not on the fly, and I don't have pickles of lists of groups.
import pandas as pd
import numpy as np

np.random.seed(0)

df = pd.DataFrame({"a"     : np.random.normal(0, 1, 100),
                   "index1": np.random.randint(0, 5, 100),
                   "index2": np.random.randint(0, 5, 100)})

grouped_df = df.groupby(["index1", "index2"])

# convert dataframe to array for e.g. deep learning but never do this operation again (too slow)
array = np.array([group["a"].values for _, group in grouped_df])

# fetch the same sample from the array and the df
array_n = array[3] # this is trivial

# how can I do this in my df?
# grouped_df[3] obviously doesn't work.

A:

You can get a list of group keys using grouped_df.groups attribute, which returns a dict. Then get the keys of that dict using the keys() function:
In [27]: grouped_df.groups.keys()
Out[27]: dict_keys([(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (2, 0), (2, 1), (2, 3), (2, 4), (3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (4, 0), (4, 1), (4, 2), (4, 3), (4, 4)])

(These keys have values of (index1, index2) corresponding to the index1/index2 value for that group.)
The list comprehension you use to populate array will iterate through these group keys in order, so element 4 of the resulting array corresponds to the fourth key:
In [28]: list(grouped_df.groups.keys())[3]
Out[28]: (0, 3)

(which indicates index1=0 and index2=3). Now you can assign that key to a variable and get the group corresponding to that key using the grouped_df.get_group() method:
In [29]: my_key = list(grouped_df.groups.keys())[3]

In [30]: grouped_df.get_group(my_key)
Out[30]:
           a  index1  index2
0   1.764052       0       3
14  0.443863       0       3
56  0.066517       0       3
58 -0.634322       0       3
65 -0.401781       0       3
69  0.051945       0       3

Finally, grab column "a" from the result, and get the values the same way you did with the list comprehension:
In [31]: grouped_df.get_group(my_key)['a'].values
Out[31]:
array([ 1.76405235,  0.44386323,  0.06651722, -0.63432209, -0.40178094,
        0.0519454 ])

Or, as a somewhat messy one-liner,
In [32]: grouped_df.get_group(list(grouped_df.groups.keys())[3])['a'].values
Out[32]:
array([ 1.76405235,  0.44386323,  0.06651722, -0.63432209, -0.40178094,
        0.0519454 ])

