Панды: выберите первые 3 строки с разными значениями в определенном столбце

Question

Я работаю с фреймом данных Pandas следующим образом:

  col1 col2 id  name
0   1   2   1   harry
1   2   3   1   harry
2   3   4   2   harry
3   4   5   3   harry
4   5   6   4   harry
5   1   9   6   jane
6   2   10  6   jane
7   3   11  7   jane
8   4   12  7   jane
9   5   12  8   jane

Я хотел бы создать новый фрейм данных, в котором для каждого name я беру первые 3 строки с уникальным id values.

Например, для harry это будут строки 0, 2 и 3. Строка 1 исключается, поскольку она имеет тот же идентификатор, что и строка 0.

Правильныйвыходные данные для моего примера dataframe:

  col1 col2 id  name
0   1   2   1   harry
2   3   4   2   harry
3   4   5   3   harry
5   1   9   6   jane
7   3   11  7   jane
9   5   12  8   jane

Пример dataframe задается этим кодом:

example = pd.DataFrame({"col1":[1,2,3,4,5, 1,2,3,4,5], "id":[1,1,2,3,4, 6, 6, 7, 7, 8],
                        "col2":[2,3,4,5,6, 9, 10, 11, 12, 12],
                        "name":["harry", "harry", "harry", "harry", "harry", "jane",
                               "jane","jane","jane","jane",]})

Этот код работает, но очень уродлив и не векторизован:

result_df = pd.DataFrame(columns=example.columns)
names_to_ids = {}
for i, row in example.iterrows():
    curr_name = row["name"]
    curr_id = row["id"]
    print curr_name, curr_id
    if curr_name not in names_to_ids:
        result_df = result_df.append(row)
        names_to_ids[curr_name] = [curr_id]
    elif len(names_to_ids[curr_name]) < 3 and curr_id not in names_to_ids[curr_name]:
        result_df = result_df.append(row)
        names_to_ids[curr_name].append(curr_id)

Answer 1

Другой способ - использовать двойные groupby и head

In [183]: df.groupby(['name', 'id']).head(1).groupby('name').head(3)
Out[183]:
   col1  col2  id   name
0     1     2   1  harry
2     3     4   2  harry
3     4     5   3  harry
5     1     9   6   jane
7     3    11   7   jane
9     5    12   8   jane

Answer 2

Используя drop_duplicates, а затем GroupBy + cumcount:

res = df.drop_duplicates(['id', 'name'])
res = res.loc[res.groupby('name').cumcount().lt(3)]

print(res)

   col1  col2  id   name
0     1     2   1  harry
2     3     4   2  harry
3     4     5   3  harry
5     1     9   6   jane
7     3    11   7   jane
9     5    12   8   jane

Answer 3

Использование drop_duplicates и head:

df.drop_duplicates(['id', 'name']).groupby('name').head(3)

   col1  col2  id   name
0     1     2   1  harry
2     3     4   2  harry
3     4     5   3  harry
5     1     9   6   jane
7     3    11   7   jane
9     5    12   8   jane