python - 使用 iterrows 选择 pandas 数据框中的下 N 行

Question

我需要使用 iterrows 在 pandas Dataframe 中每次选择 N 行。像这样:

def func():
    selected = []
    for i in range(N):
        selected.append(next(dataframe.iterrows()))

    yield selected

但是这样做 selected 有 N 个相等的元素。每次我调用 func 我总是得到相同的结果(数据帧的第一个元素)。

如果数据框是:

   A  B  C
0  5  8  2
1  1  2  3
2  4  5  6
3  7  8  9
4  0  1  2
5  3  4  5
6  7  8  6
7  1  2  3

我要获取的是:

N = 3
selected = [ [5,8,2], [1,2,3], [4,5,6] ] 
then, calling again the function,
selected = [ [7,8,9], [0,1,2], [3,4,5] ] 
then,
selected = [ [7,8,6], [1,2,3], [5,8,2] ] 

Answer 1

不需要.iterrows()，而是使用切片:

def flow_from_df(dataframe: pd.DataFrame, chunk_size: int = 10):
    for start_row in range(0, dataframe.shape[0], chunk_size):
        end_row  = min(start_row + chunk_size, dataframe.shape[0])
        yield dataframe.iloc[start_row:end_row, :]

使用方法:

get_chunk = flow_from_df(dataframe)
chunk1 = next(get_chunk)
chunk2 = next(get_chunk)

或者不使用生成器:

def get_chunk(dataframe: pd.DataFrame, chunk_size: int, start_row: int = 0) -> pd.DataFrame:
    end_row  = min(start_row + chunk_size, dataframe.shape[0])

    return dataframe.iloc[start_row:end_row, :]