如何从pandas DataFrame中提取指定2D点的逐帧坐标数组

以下提供两种符合Python风格的实现，分别适用不同场景：

1. 易读优先实现（适合小数据集、代码可读性要求高的场景） #

import pandas as pd
import numpy as np

def extract_points(df, indices):
    '''
    传入DataFrame和点索引列表
    返回对应点的坐标列表
    '''
    return np.array([
        [tuple(coord) for coord in zip(df[f'point_{i}_x'], df[f'point_{i}_y'])]
        for i in indices
    ])

说明 #

• 用列表推导式遍历所有目标点索引，逻辑直观易懂
• 用zip直接将同一帧的x、y坐标配对为元组，完全匹配示例输出的结构要求
• 外层直接包裹np.array转换为目标类数组结构

2. 性能优先实现（适合大数据集、索引数量多的场景） #

def extract_points(df, indices):
    '''
    传入DataFrame和点索引列表
    返回对应点的坐标列表
    '''
    # 构造需要提取的坐标列名
    target_cols = [f'point_{i}_{ax}' for i in indices for ax in ('x', 'y')]
    # 向量化提取+维度转换，无Python层循环
    coord_arr = df[target_cols].values.reshape(len(df), len(indices), 2).transpose(1, 0, 2)
    # 如果需要严格返回元组结构，可加下面一行
    # coord_arr = np.array([[tuple(p) for p in point] for point in coord_arr])
    return coord_arr

说明 #

• 全程使用pandas+numpy向量化操作，避免Python层循环，数据量越大性能优势越明显
• 通过reshape和transpose直接调整维度到目标结构，代码简洁高效

测试验证 #

用你提供的示例数据集测试：

# 构造示例数据集
data = {
    'frame': [1,2,3],
    'point_1_x': [0,2,8],
    'point_2_x': [1,3,2],
    'point_3_x': [1,5,3],
    'point_1_y': [2,1,4],
    'point_2_y': [3,2,5],
    'point_3_y': [1,3,6]
}
example_dataset = pd.DataFrame(data)

# 调用函数测试
print(extract_points(example_dataset, [1,2]))

输出完全匹配要求：

[[[0 2]
  [2 1]
  [8 4]]

 [[1 3]
  [3 2]
  [2 5]]]

内容的提问来源于stack exchange，提问作者Nourless