python - 如何为熊猫数据框中的每一行创建坐标列表？

Question

我有一个数据框，其中包含帧序列中几个 2d 点的坐标。看起来像

frame point_1_x point_2_x point_3_x point_1_y point_2_y point_3_y
1          0         1         1         2         3         1
2          2         3         5         1         2         3
3          8         2         3         4         5         6

我想以类似数组的结构形式提取给定点索引的坐标。例如：

def extract_points(df, indices):
    '''
        Takes dataframe and indices of points
        returns list of coordinates of points
    '''
extract_points(example_dataset, [1,2])
output: np.array([[(0, 2), (2,1), (8, 4)], [(1,3), (3,2), (2,5)]])

我怎样才能使用numpy和以pythonic方式做到这一点pandas？

标签： pythonpandas

Answer 1

重命名列，然后 groupby

df.columns = df.columns.str[:-2]
arr = df.stack().groupby(level=[1,0]).agg(tuple).values

array([(0, 2), (2, 1), (8, 4), (1, 3), (3, 2), (2, 5), (1, 1), (5, 3),
       (3, 6)], dtype=object)