Repository: pandas-dev/pandas

Issue:
ENH: enumerate groups
Sometimes it's handy to have access to a distinct integer for each group.  For example, using the (internal) grouper:

```
>>> df = pd.DataFrame({"a": list("xyyzxy"), "b": list("ab"*3), "c": range(6)})
>>> df["group_id"] = df.groupby(["a","b"]).grouper.group_info[0]
>>> df
   a  b  c  group_id
0  x  a  0         0
1  y  b  1         2
2  y  a  2         1
3  z  b  3         3
4  x  a  4         0
5  y  b  5         2
```

This can be achieved in a number of ways but none of them are particularly elegant, esp. if we're grouping on multiple keys and/or Series.  Accordingly, after a brief discussion on gitter, I propose a new method `transform("enumerate")` which returns a Series of integers from 0 to ngroups-1 matching the order the groups will be iterated in.  In other words, we'll simply be applying the following map:

```
>>> m = {k: i for i, (k,g) in enumerate(df.groupby(["a","b"]))}
>>> m
{('x', 'a'): 0, ('y', 'b'): 2, ('y', 'a'): 1, ('z', 'b'): 3}
```

(Note this is only to shows the desired behaviour, and wouldn't be how it'd be implemented!)



Discussion:
can you show an example of its utility!

also to note that this is really only useful as a `.transform` method (a reduction is kind of silly as its just the `range(len(df.groupby(...)))`)

Note that this is essentially exactly the same information provided by `pandas.factorize`:

```
In [1]: import pandas as pd

In [2]: pd.factorize(['a', 'a', 'b', 'c'])
Out[2]: (array([0, 0, 1, 2]), array(['a', 'b', 'c'], dtype=object))
```

I couldn't think of a clean way to get factorize to handle the same inputs as groupby, though (both the multiple-series case and the mixed column-name/list input.)  Might have missed something obvious, of course, as is my wont.  But if I needed to write a few lines to get it to work, then those lines would more naturally fit as a groupby method, or so it seemed to me.

As I went to implement this, I started to wonder if it doesn't make more sense to use `df.groupby("a").enumerate()` instead of `df.groupby("a").transform("enumerate")`, to be parallel with `df.groupby("a").cumcount()`, instead of `df.groupby("a").transform("cumcount")` (which doesn't work.)  This would give us something like

```
>>> df = pd.DataFrame({"A": [1,2,2,2,1]})
>>> df["group_id"] = df.groupby("A").enumerate()
>>> df["group_index"] = df.groupby("A").cumcount()
>>> df
   A  group_id  group_index
0  1         0            0
1  2         1            0
2  2         1            1
3  2         1            2
4  1         0            1
```

that looks reasonable 


Resolving patch:
diff --git a/pandas/core/groupby.py b/pandas/core/groupby.py
--- a/pandas/core/groupby.py
+++ b/pandas/core/groupby.py
@@ -150,7 +150,7 @@
     'last', 'first',
     'head', 'tail', 'median',
     'mean', 'sum', 'min', 'max',
-    'cumcount',
+    'cumcount', 'ngroup',
     'resample',
     'rank', 'quantile',
     'fillna',
@@ -1437,6 +1437,75 @@ def nth(self, n, dropna=None):
 
         return result
 
+    @Substitution(name='groupby')
+    @Appender(_doc_template)
+    def ngroup(self, ascending=True):
+        """
+        Number each group from 0 to the number of groups - 1.
+
+        This is the enumerative complement of cumcount.  Note that the
+        numbers given to the groups match the order in which the groups
+        would be seen when iterating over the groupby object, not the
+        order they are first observed.
+
+        .. versionadded:: 0.20.2
+
+        Parameters
+        ----------
+        ascending : bool, default True
+            If False, number in reverse, from number of group - 1 to 0.
+
+        Examples
+        --------
+
+        >>> df = pd.DataFrame({"A": list("aaabba")})
+        >>> df
+           A
+        0  a
+        1  a
+        2  a
+        3  b
+        4  b
+        5  a
+        >>> df.groupby('A').ngroup()
+        0    0
+        1    0
+        2    0
+        3    1
+        4    1
+        5    0
+        dtype: int64
+        >>> df.groupby('A').ngroup(ascending=False)
+        0    1
+        1    1
+        2    1
+        3    0
+        4    0
+        5    1
+        dtype: int64
+        >>> df.groupby(["A", [1,1,2,3,2,1]]).ngroup()
+        0    0
+        1    0
+        2    1
+        3    3
+        4    2
+        5    0
+        dtype: int64
+
+        See also
+        --------
+        .cumcount : Number the rows in each group.
+
+        """
+
+        self._set_group_selection()
+
+        index = self._selected_obj.index
+        result = Series(self.grouper.group_info[0], index)
+        if not ascending:
+            result = self.ngroups - 1 - result
+        return result
+
     @Substitution(name='groupby')
     @Appender(_doc_template)
     def cumcount(self, ascending=True):
@@ -1481,6 +1550,10 @@ def cumcount(self, ascending=True):
         4    0
         5    0
         dtype: int64
+
+        See also
+        --------
+        .ngroup : Number the groups themselves.
         """
 
         self._set_group_selection()


Test patch:
diff --git a/pandas/tests/groupby/test_counting.py b/pandas/tests/groupby/test_counting.py
new file mode 100644
--- /dev/null
+++ b/pandas/tests/groupby/test_counting.py
@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+from __future__ import print_function
+
+import numpy as np
+
+from pandas import (DataFrame, Series, MultiIndex)
+from pandas.util.testing import assert_series_equal
+from pandas.compat import (range, product as cart_product)
+
+
+class TestCounting(object):
+
+    def test_cumcount(self):
+        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'])
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0, 1, 2, 0, 3])
+
+        assert_series_equal(expected, g.cumcount())
+        assert_series_equal(expected, sg.cumcount())
+
+    def test_cumcount_empty(self):
+        ge = DataFrame().groupby(level=0)
+        se = Series().groupby(level=0)
+
+        # edge case, as this is usually considered float
+        e = Series(dtype='int64')
+
+        assert_series_equal(e, ge.cumcount())
+        assert_series_equal(e, se.cumcount())
+
+    def test_cumcount_dupe_index(self):
+        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
+                       index=[0] * 5)
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0, 1, 2, 0, 3], index=[0] * 5)
+
+        assert_series_equal(expected, g.cumcount())
+        assert_series_equal(expected, sg.cumcount())
+
+    def test_cumcount_mi(self):
+        mi = MultiIndex.from_tuples([[0, 1], [1, 2], [2, 2], [2, 2], [1, 0]])
+        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
+                       index=mi)
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0, 1, 2, 0, 3], index=mi)
+
+        assert_series_equal(expected, g.cumcount())
+        assert_series_equal(expected, sg.cumcount())
+
+    def test_cumcount_groupby_not_col(self):
+        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
+                       index=[0] * 5)
+        g = df.groupby([0, 0, 0, 1, 0])
+        sg = g.A
+
+        expected = Series([0, 1, 2, 0, 3], index=[0] * 5)
+
+        assert_series_equal(expected, g.cumcount())
+        assert_series_equal(expected, sg.cumcount())
+
+    def test_ngroup(self):
+        df = DataFrame({'A': list('aaaba')})
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0, 0, 0, 1, 0])
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_distinct(self):
+        df = DataFrame({'A': list('abcde')})
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series(range(5), dtype='int64')
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_one_group(self):
+        df = DataFrame({'A': [0] * 5})
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0] * 5)
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_empty(self):
+        ge = DataFrame().groupby(level=0)
+        se = Series().groupby(level=0)
+
+        # edge case, as this is usually considered float
+        e = Series(dtype='int64')
+
+        assert_series_equal(e, ge.ngroup())
+        assert_series_equal(e, se.ngroup())
+
+    def test_ngroup_series_matches_frame(self):
+        df = DataFrame({'A': list('aaaba')})
+        s = Series(list('aaaba'))
+
+        assert_series_equal(df.groupby(s).ngroup(),
+                            s.groupby(s).ngroup())
+
+    def test_ngroup_dupe_index(self):
+        df = DataFrame({'A': list('aaaba')}, index=[0] * 5)
+        g = df.groupby('A')
+        sg = g.A
+
+        expected = Series([0, 0, 0, 1, 0], index=[0] * 5)
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_mi(self):
+        mi = MultiIndex.from_tuples([[0, 1], [1, 2], [2, 2], [2, 2], [1, 0]])
+        df = DataFrame({'A': list('aaaba')}, index=mi)
+        g = df.groupby('A')
+        sg = g.A
+        expected = Series([0, 0, 0, 1, 0], index=mi)
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_groupby_not_col(self):
+        df = DataFrame({'A': list('aaaba')}, index=[0] * 5)
+        g = df.groupby([0, 0, 0, 1, 0])
+        sg = g.A
+
+        expected = Series([0, 0, 0, 1, 0], index=[0] * 5)
+
+        assert_series_equal(expected, g.ngroup())
+        assert_series_equal(expected, sg.ngroup())
+
+    def test_ngroup_descending(self):
+        df = DataFrame(['a', 'a', 'b', 'a', 'b'], columns=['A'])
+        g = df.groupby(['A'])
+
+        ascending = Series([0, 0, 1, 0, 1])
+        descending = Series([1, 1, 0, 1, 0])
+
+        assert_series_equal(descending, (g.ngroups - 1) - ascending)
+        assert_series_equal(ascending, g.ngroup(ascending=True))
+        assert_series_equal(descending, g.ngroup(ascending=False))
+
+    def test_ngroup_matches_cumcount(self):
+        # verify one manually-worked out case works
+        df = DataFrame([['a', 'x'], ['a', 'y'], ['b', 'x'],
+                        ['a', 'x'], ['b', 'y']], columns=['A', 'X'])
+        g = df.groupby(['A', 'X'])
+        g_ngroup = g.ngroup()
+        g_cumcount = g.cumcount()
+        expected_ngroup = Series([0, 1, 2, 0, 3])
+        expected_cumcount = Series([0, 0, 0, 1, 0])
+
+        assert_series_equal(g_ngroup, expected_ngroup)
+        assert_series_equal(g_cumcount, expected_cumcount)
+
+    def test_ngroup_cumcount_pair(self):
+        # brute force comparison for all small series
+        for p in cart_product(range(3), repeat=4):
+            df = DataFrame({'a': p})
+            g = df.groupby(['a'])
+
+            order = sorted(set(p))
+            ngroupd = [order.index(val) for val in p]
+            cumcounted = [p[:i].count(val) for i, val in enumerate(p)]
+
+            assert_series_equal(g.ngroup(), Series(ngroupd))
+            assert_series_equal(g.cumcount(), Series(cumcounted))
+
+    def test_ngroup_respects_groupby_order(self):
+        np.random.seed(0)
+        df = DataFrame({'a': np.random.choice(list('abcdef'), 100)})
+        for sort_flag in (False, True):
+            g = df.groupby(['a'], sort=sort_flag)
+            df['group_id'] = -1
+            df['group_index'] = -1
+
+            for i, (_, group) in enumerate(g):
+                df.loc[group.index, 'group_id'] = i
+                for j, ind in enumerate(group.index):
+                    df.loc[ind, 'group_index'] = j
+
+            assert_series_equal(Series(df['group_id'].values),
+                                g.ngroup())
+            assert_series_equal(Series(df['group_index'].values),
+                                g.cumcount())
diff --git a/pandas/tests/groupby/test_groupby.py b/pandas/tests/groupby/test_groupby.py
--- a/pandas/tests/groupby/test_groupby.py
+++ b/pandas/tests/groupby/test_groupby.py
@@ -3399,60 +3399,6 @@ def test_groupby_with_small_elem(self):
         res = grouped.get_group((pd.Timestamp('2014-08-31'), 'start'))
         tm.assert_frame_equal(res, df.iloc[[2], :])
 
-    def test_cumcount(self):
-        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'])
-        g = df.groupby('A')
-        sg = g.A
-
-        expected = Series([0, 1, 2, 0, 3])
-
-        assert_series_equal(expected, g.cumcount())
-        assert_series_equal(expected, sg.cumcount())
-
-    def test_cumcount_empty(self):
-        ge = DataFrame().groupby(level=0)
-        se = Series().groupby(level=0)
-
-        # edge case, as this is usually considered float
-        e = Series(dtype='int64')
-
-        assert_series_equal(e, ge.cumcount())
-        assert_series_equal(e, se.cumcount())
-
-    def test_cumcount_dupe_index(self):
-        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
-                       index=[0] * 5)
-        g = df.groupby('A')
-        sg = g.A
-
-        expected = Series([0, 1, 2, 0, 3], index=[0] * 5)
-
-        assert_series_equal(expected, g.cumcount())
-        assert_series_equal(expected, sg.cumcount())
-
-    def test_cumcount_mi(self):
-        mi = MultiIndex.from_tuples([[0, 1], [1, 2], [2, 2], [2, 2], [1, 0]])
-        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
-                       index=mi)
-        g = df.groupby('A')
-        sg = g.A
-
-        expected = Series([0, 1, 2, 0, 3], index=mi)
-
-        assert_series_equal(expected, g.cumcount())
-        assert_series_equal(expected, sg.cumcount())
-
-    def test_cumcount_groupby_not_col(self):
-        df = DataFrame([['a'], ['a'], ['a'], ['b'], ['a']], columns=['A'],
-                       index=[0] * 5)
-        g = df.groupby([0, 0, 0, 1, 0])
-        sg = g.A
-
-        expected = Series([0, 1, 2, 0, 3], index=[0] * 5)
-
-        assert_series_equal(expected, g.cumcount())
-        assert_series_equal(expected, sg.cumcount())
-
     def test_fill_constistency(self):
 
         # GH9221
diff --git a/pandas/tests/groupby/test_whitelist.py b/pandas/tests/groupby/test_whitelist.py
--- a/pandas/tests/groupby/test_whitelist.py
+++ b/pandas/tests/groupby/test_whitelist.py
@@ -24,6 +24,7 @@
     'head',
     'tail',
     'cumcount',
+    'ngroup',
     'resample',
     'rank',
     'quantile',
@@ -61,6 +62,7 @@
     'head',
     'tail',
     'cumcount',
+    'ngroup',
     'resample',
     'rank',
     'quantile',
@@ -237,7 +239,7 @@ def test_tab_completion(mframe):
          'prod', 'size', 'std', 'sum', 'transform', 'var', 'sem', 'count',
          'nunique', 'head', 'describe', 'cummax', 'quantile',
          'rank', 'cumprod', 'tail', 'resample', 'cummin', 'fillna',
-         'cumsum', 'cumcount', 'all', 'shift', 'skew',
+         'cumsum', 'cumcount', 'ngroup', 'all', 'shift', 'skew',
          'take', 'tshift', 'pct_change', 'any', 'mad', 'corr', 'corrwith',
          'cov', 'dtypes', 'ndim', 'diff', 'idxmax', 'idxmin',
          'ffill', 'bfill', 'pad', 'backfill', 'rolling', 'expanding'])