def test_nth(self):
    df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
    g = df.groupby('A')
    assert_frame_equal(g.nth(0), df.iloc[[0, 2]].set_index('A'))
    assert_frame_equal(g.nth(1), df.iloc[[1]].set_index('A'))
    assert_frame_equal(g.nth(2), df.loc[[]].set_index('A'))
    assert_frame_equal(g.nth((- 1)), df.iloc[[1, 2]].set_index('A'))
    assert_frame_equal(g.nth((- 2)), df.iloc[[0]].set_index('A'))
    assert_frame_equal(g.nth((- 3)), df.loc[[]].set_index('A'))
    assert_series_equal(g.B.nth(0), df.set_index('A').B.iloc[[0, 2]])
    assert_series_equal(g.B.nth(1), df.set_index('A').B.iloc[[1]])
    assert_frame_equal(g[['B']].nth(0), df.loc[([0, 2], ['A', 'B'])].set_index('A'))
    exp = df.set_index('A')
    assert_frame_equal(g.nth(0, dropna='any'), exp.iloc[[1, 2]])
    assert_frame_equal(g.nth((- 1), dropna='any'), exp.iloc[[1, 2]])
    exp['B'] = np.nan
    assert_frame_equal(g.nth(7, dropna='any'), exp.iloc[[1, 2]])
    assert_frame_equal(g.nth(2, dropna='any'), exp.iloc[[1, 2]])
    df = DataFrame({'color': {0: 'green', 1: 'green', 2: 'red', 3: 'red', 4: 'red'}, 'food': {0: 'ham', 1: 'eggs', 2: 'eggs', 3: 'ham', 4: 'pork'}, 'two': {0: 1.5456590000000001, 1: (- 0.070345), 2: (- 2.400454), 3: 0.46206, 4: 0.523508}, 'one': {0: 0.565738, 1: (- 0.9742360000000001), 2: 1.033801, 3: (- 0.785435), 4: 0.704228}}).set_index(['color', 'food'])
    result = df.groupby(level=0, as_index=False).nth(2)
    expected = df.iloc[[(- 1)]]
    assert_frame_equal(result, expected)
    result = df.groupby(level=0, as_index=False).nth(3)
    expected = df.loc[[]]
    assert_frame_equal(result, expected)
    df = DataFrame(np.random.randint(1, 10, (100, 2)), dtype='int64')
    s = df[1]
    g = df[0]
    expected = s.groupby(g).first()
    expected2 = s.groupby(g).apply((lambda x: x.iloc[0]))
    assert_series_equal(expected2, expected, check_names=False)
    assert expected.name, 0
    assert (expected.name == 1)
    v = s[(g == 1)].iloc[0]
    assert (expected.iloc[0] == v)
    assert (expected2.iloc[0] == v)
    expected = s.groupby(g, sort=False).first()
    result = s.groupby(g, sort=False).nth(0, dropna='all')
    assert_series_equal(result, expected)
    df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
    g = df.groupby('A')
    result = g.B.nth(0, dropna=True)
    expected = g.B.first()
    assert_series_equal(result, expected)
    df = DataFrame([[1, np.nan], [1, 3], [1, 4], [5, 6], [5, 7]], columns=['A', 'B'])
    g = df.groupby('A')
    assert_frame_equal(g.nth(0), df.iloc[[0, 3]].set_index('A'))
    assert_frame_equal(g.nth([0]), df.iloc[[0, 3]].set_index('A'))
    assert_frame_equal(g.nth([0, 1]), df.iloc[[0, 1, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, (- 1)]), df.iloc[[0, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, 1, 2]), df.iloc[[0, 1, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, 1, (- 1)]), df.iloc[[0, 1, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([2]), df.iloc[[2]].set_index('A'))
    assert_frame_equal(g.nth([3, 4]), df.loc[[]].set_index('A'))
    business_dates = pd.date_range(start='4/1/2014', end='6/30/2014', freq='B')
    df = DataFrame(1, index=business_dates, columns=['a', 'b'])
    key = (df.index.year, df.index.month)
    result = df.groupby(key, as_index=False).nth([0, 3, (- 2), (- 1)])
    expected_dates = pd.to_datetime(['2014/4/1', '2014/4/4', '2014/4/29', '2014/4/30', '2014/5/1', '2014/5/6', '2014/5/29', '2014/5/30', '2014/6/2', '2014/6/5', '2014/6/27', '2014/6/30'])
    expected = DataFrame(1, columns=['a', 'b'], index=expected_dates)
    assert_frame_equal(result, expected)
def test_nth(self):
    df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
    g = df.groupby('A')
    assert_frame_equal(g.nth(0), df.iloc[[0, 2]].set_index('A'))
    assert_frame_equal(g.nth(1), df.iloc[[1]].set_index('A'))
    assert_frame_equal(g.nth(2), df.loc[[]].set_index('A'))
    assert_frame_equal(g.nth((- 1)), df.iloc[[1, 2]].set_index('A'))
    assert_frame_equal(g.nth((- 2)), df.iloc[[0]].set_index('A'))
    assert_frame_equal(g.nth((- 3)), df.loc[[]].set_index('A'))
    assert_series_equal(g.B.nth(0), df.set_index('A').B.iloc[[0, 2]])
    assert_series_equal(g.B.nth(1), df.set_index('A').B.iloc[[1]])
    assert_frame_equal(g[['B']].nth(0), df.loc[([0, 2], ['A', 'B'])].set_index('A'))
    exp = df.set_index('A')
    assert_frame_equal(g.nth(0, dropna='any'), exp.iloc[[1, 2]])
    assert_frame_equal(g.nth((- 1), dropna='any'), exp.iloc[[1, 2]])
    exp['B'] = np.nan
    assert_frame_equal(g.nth(7, dropna='any'), exp.iloc[[1, 2]])
    assert_frame_equal(g.nth(2, dropna='any'), exp.iloc[[1, 2]])
    df = DataFrame({'color': {0: 'green', 1: 'green', 2: 'red', 3: 'red', 4: 'red'}, 'food': {0: 'ham', 1: 'eggs', 2: 'eggs', 3: 'ham', 4: 'pork'}, 'two': {0: 1.5456590000000001, 1: (- 0.070345), 2: (- 2.400454), 3: 0.46206, 4: 0.523508}, 'one': {0: 0.565738, 1: (- 0.9742360000000001), 2: 1.033801, 3: (- 0.785435), 4: 0.704228}}).set_index(['color', 'food'])
    result = df.groupby(level=0, as_index=False).nth(2)
    expected = df.iloc[[(- 1)]]
    assert_frame_equal(result, expected)
    result = df.groupby(level=0, as_index=False).nth(3)
    expected = df.loc[[]]
    assert_frame_equal(result, expected)
    df = DataFrame(np.random.randint(1, 10, (100, 2)), dtype='int64')
    s = df[1]
    g = df[0]
    expected = s.groupby(g).first()
    expected2 = s.groupby(g).apply((lambda x: x.iloc[0]))
    assert_series_equal(expected2, expected, check_names=False)
    assert expected.name, 0
    assert (expected.name == 1)
    v = s[(g == 1)].iloc[0]
    assert (expected.iloc[0] == v)
    assert (expected2.iloc[0] == v)
    expected = s.groupby(g, sort=False).first()
    result = s.groupby(g, sort=False).nth(0, dropna='all')
    assert_series_equal(result, expected)
    df = DataFrame([[1, np.nan], [1, 4], [5, 6]], columns=['A', 'B'])
    g = df.groupby('A')
    result = g.B.nth(0, dropna=True)
    expected = g.B.first()
    assert_series_equal(result, expected)
    df = DataFrame([[1, np.nan], [1, 3], [1, 4], [5, 6], [5, 7]], columns=['A', 'B'])
    g = df.groupby('A')
    assert_frame_equal(g.nth(0), df.iloc[[0, 3]].set_index('A'))
    assert_frame_equal(g.nth([0]), df.iloc[[0, 3]].set_index('A'))
    assert_frame_equal(g.nth([0, 1]), df.iloc[[0, 1, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, (- 1)]), df.iloc[[0, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, 1, 2]), df.iloc[[0, 1, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([0, 1, (- 1)]), df.iloc[[0, 1, 2, 3, 4]].set_index('A'))
    assert_frame_equal(g.nth([2]), df.iloc[[2]].set_index('A'))
    assert_frame_equal(g.nth([3, 4]), df.loc[[]].set_index('A'))
    business_dates = pd.date_range(start='4/1/2014', end='6/30/2014', freq='B')
    df = DataFrame(1, index=business_dates, columns=['a', 'b'])
    key = (df.index.year, df.index.month)
    result = df.groupby(key, as_index=False).nth([0, 3, (- 2), (- 1)])
    expected_dates = pd.to_datetime(['2014/4/1', '2014/4/4', '2014/4/29', '2014/4/30', '2014/5/1', '2014/5/6', '2014/5/29', '2014/5/30', '2014/6/2', '2014/6/5', '2014/6/27', '2014/6/30'])
    expected = DataFrame(1, columns=['a', 'b'], index=expected_dates)
    assert_frame_equal(result, expected)