Pandas multiindex dataframe set first row in a column to 0

Question

I am having some trouble working on grouped objects in pandas. Specifically, I want to be able to set the first row in a column to 0 while keeping other rows unchanged.

For example:

df = pd.DataFrame({'A': ['foo', 'bar', 'baz'] * 2,
                        'B': rand.randn(6),
                        'C': rand.rand(6) > .5})

gives me

    A         B      C
0  foo  1.624345  False
1  bar -0.611756   True
2  baz -0.528172  False
3  foo -1.072969   True
4  bar  0.865408  False
5  baz -2.301539   True

and I group them by A and sort them by B:

f = lambda x: x.sort('B', ascending=True)
sort_df = df.groupby('A',sort=False).apply(f)

to get this:

         A         B      C
    A                          
foo 3  foo -1.072969   True
    0  foo  1.624345  False
bar 1  bar -0.611756   True
    4  bar  0.865408  False
baz 5  baz -2.301539   True
    2  baz -0.528172  False

Now that I have the groups, I want to be able to set the first element in each group to 0. How do I do that?

Something like this would work, but I want a more optimized way of doing it:

for group in sort_df.groupby(level=0).groups:
    sort_df.loc[(group,sort_df.loc[group].index[0]),'B']=0

Any help would be appreciated, thanks!

Answer 1

Here's the vectorized way to do this. Will be MUCH faster

In [26]: pd.set_option('max_rows',10)

Create a frame with a 2-level multi-index, sort by A (arbitrary here)

In [27]: df = DataFrame(dict(A = np.random.randint(0,100,size=N),B=np.random.randint(0,100,size=N),C=np.random.randn(N))).sort(columns=['A'])

In [28]: df
Out[28]: 
        A   B         C
61474   0  40 -0.731163
82386   0  18 -1.316136
63372   0  28  0.112819
49666   0  13 -0.649491
31631   0  89 -0.835208
...    ..  ..       ...
42178  99  28 -0.029800
59529  99  31 -0.733588
13503  99  60  0.672754
20961  99  18  0.252714
31882  99  22  0.083340

[100000 rows x 3 columns]

Reset the index to capture the index value. Find the first values according to B

In [29]: grouped = df.reset_index().groupby('B').first()

In [30]: grouped
Out[30]: 
    index  A         C
B                     
0   26576  0  1.123605
1   38311  0  0.128966
2   45135  0 -0.039886
3   38434  0 -1.284028
4   82088  0 -0.747440
..    ... ..       ...
95  82620  0 -1.197625
96  63278  0 -0.625400
97  23226  0 -0.497609
98  82520  0 -0.828773
99  35902  0 -0.199752

[100 rows x 3 columns]

This gives you an indexer into the frame.

In [31]: df.loc[grouped['index']] = 0

In [32]: df
Out[32]: 
        A   B         C
61474   0   0  0.000000
82386   0   0  0.000000
63372   0   0  0.000000
49666   0   0  0.000000
31631   0   0  0.000000
...    ..  ..       ...
42178  99  28 -0.029800
59529  99  31 -0.733588
13503  99  60  0.672754
20961  99  18  0.252714
31882  99  22  0.083340

[100000 rows x 3 columns]

If you want

In [33]: df.sort_index()
Out[33]: 
        A   B         C
0      40  56 -1.223941
1      24  77 -0.039775
2       7  83  0.741013
3      48  38 -1.795053
4      62  15 -2.734968
...    ..  ..       ...
99995  20  25 -0.286300
99996  27  21 -0.120430
99997   0   4  0.607524
99998  38  31  0.717069
99999  33  63 -0.226888

[100000 rows x 3 columns]

This method

In [34]: %timeit df.loc[grouped['index']] = 0
100 loops, best of 3: 7.33 ms per loop

Your original

In [37]: %timeit df.groupby('A',sort=False).apply(f)
10 loops, best of 3: 109 ms per loop

If you have a lot more groups this perf differnce will widen.