问使用python选择csv文件中每个组的前k项

EN

import pandas as pd

df = pd.read_csv(f_dir)
fd = df.groupby('category').head(2)
print fd

catsSeen = {}
returnList = []
for row in table:
    if row[1] not in catsSeen:
        catsSeen[row[1]] = 1
        returnList.append(row)
    else:
        if catsSeen[row[1]] < 2:
            returnList.append(row)
            catsSeen[row[1]] += 1

print(returnList)

from io import StringIO
from csv import DictReader, DictWriter
from itertools import groupby

txt = '''ID.   category
 1.     apple
 2.     apple
 3.     apple
 4.     apple
 5.     orange
 6.     orange
 7.     orange
 8.     banana
 9.     banana
 10.    banana'''

k = 2

with StringIO(txt) as in_file, StringIO() as out_file:

    data = DictReader(in_file, delimiter=' ', skipinitialspace=True)
    out = DictWriter(out_file, delimiter=' ', fieldnames=data.fieldnames)
    out.writeheader()

    for key, items in groupby(data, key=lambda x: x['category']):
        for item, _ in zip(items, range(k)):
            out.writerow(item)

    print(out_file.getvalue())  # this is just to print the result...

ID. category
1. apple
2. apple
5. orange
6. orange
8. banana
9. banana