diff --git a/pandas/core/index.py b/pandas/core/index.py index 2de9b2fac9aa9..6702a21167850 100644 --- a/pandas/core/index.py +++ b/pandas/core/index.py @@ -1555,6 +1555,20 @@ def _possibly_promote(self, other): return self, other def groupby(self, to_groupby): + """ + Group the index labels by a given array of values. + + Parameters + ---------- + to_groupby : array + Values used to determine the groups. + + Returns + ------- + groups : dict + {group name -> group labels} + + """ return self._groupby(self.values, _values_from_object(to_groupby)) def map(self, mapper): diff --git a/pandas/tests/test_index.py b/pandas/tests/test_index.py index adb5e7d07fbe6..cca8324b42b93 100644 --- a/pandas/tests/test_index.py +++ b/pandas/tests/test_index.py @@ -1137,6 +1137,11 @@ def get_reindex_type(target): self.assertEqual(reindexed.levels[0].dtype.type, np.int64) self.assertEqual(reindexed.levels[1].dtype.type, np.float64) + def test_groupby(self): + idx = Index(range(5)) + groups = idx.groupby(np.array([1,1,2,2,2])) + exp = {1: [0, 1], 2: [2, 3, 4]} + tm.assert_dict_equal(groups, exp) class Numeric(Base): @@ -3524,6 +3529,17 @@ def test_reindex_lvl_preserves_type_if_target_is_empty_list_or_array(self): self.assertEqual(idx.reindex([], level=1)[0].levels[1].dtype.type, np.object_) + def test_groupby(self): + groups = self.index.groupby(np.array([1, 1, 1, 2, 2, 2])) + labels = self.index.get_values().tolist() + exp = {1: labels[:3], 2: labels[3:]} + tm.assert_dict_equal(groups, exp) + + # GH5620 + groups = self.index.groupby(self.index) + exp = dict((key, [key]) for key in self.index) + tm.assert_dict_equal(groups, exp) + def test_get_combined_index(): from pandas.core.index import _get_combined_index