Pandas: Multi-Index Subset Selection

Viewed 51

The following code will produce a dataframe that represents four players, each making multiple rolls of two dice:

mux = pd.MultiIndex.from_arrays([
    list(['Alice', 'Alice', 'Alice', 
          'Bob', 'Bob', 'Bob', 'Bob', 'Bob', 
          'Carol', 'Carol', 'Carol', 'Carol', 'Carol', 'Carol', 
          'Dan', 'Dan']),
    list('1231234512345612')
], names=['player', 'roll'])

#df = pd.DataFrame({'die1': np.random.randint(1, 7, len(mux)),
#                   'die2': np.random.randint(1, 7, len(mux))}, mux)

df = pd.DataFrame({'die1': ['3','6','2','1','1','6','5','1','3','1','4','5','3','5','5','5'],
                   'die2': ['2','5','1','4','1','6','3','5','6','4','5','1','4','4','3','6']}, mux)

This is the dataframe produced:

            die1 die2
player roll          
Alice  1       3    2
       2       6    5
       3       2    1
Bob    1       1    4
       2       1    1
       3       6    6
       4       5    3
       5       1    5
Carol  1       3    6
       2       1    4
       3       4    5
       4       5    1
       5       3    4
       6       5    4
Dan    1       5    3
       2       5    6

I am interested in the players-- not just the individual rolls, but the players-- who roll a 6 on the first die.

The following code does what I do not want, it returns the individual rolls where the first die is a 6: df[df.die1=='6']:

             die1   die2
player  roll        
Alice   2    6      5
Bob     3    6      6

How can I produce the following, by using die1, without making mention of the individual player names, which I obviously won't know a priori?

            die1 die2
player roll          
Alice  1       3    2
       2       6    5
       3       2    1
Bob    1       1    4
       2       1    1
       3       6    6
       4       5    3
       5       1    5
2 Answers

You can do groupby:

# is your data string '6' or integer 6
df[df['die1'].eq('6').groupby('player').transform('any')]

Output:

            die1 die2
player roll          
Alice  1       3    2
       2       6    5
       3       2    1
Bob    1       1    4
       2       1    1
       3       6    6
       4       5    3
       5       1    5
ll=[]
for i, j in df.groupby('player'):
    if len(j[j.die1==6])>0:
        ll.append(j)
return pd.concat(ll)
Related