Exemplo n.º 1
0
def test_dataframe_nsmallest(nelem, n):
    np.random.seed(0)
    df = DataFrame()
    df["a"] = aa = np.random.random(nelem)
    df["b"] = bb = np.random.random(nelem)
    res = df.nsmallest(n, "a")

    # Check
    inds = np.argsort(-aa)
    assert_eq(res["a"].to_array(), aa[inds][-n:][::-1])
    assert_eq(res["b"].to_array(), bb[inds][-n:][::-1])
    assert_eq(res.index.values, inds[-n:][::-1])