如何使用列名数组有效读取 pandas

Question

df = pd.DataFrame({"col_a": [1,2,3], "col_b": [5,4,0], "col_c": [9,7,6])
cols = [["col_a", "col_b"], ["col_c", "col_b"], ["col_a", "col_b"]]

预期输出：

[[1,5], [7,4], [3,0]]

我知道这可以使用列表理解来实现，寻找更有效的方法，因为我有超过一百万条记录

Answer 1

您忘记提供的列表理解：

In [27]: [row[1][col].to_list() for row, col in zip(df.iterrows(), cols)]
Out[27]: [[1, 5], [7, 4], [3, 0]]

Answer 2

我认为如果不迭代

cols

变量，你就无法做到这一点。试试这个-

[df.loc[i,j].tolist() for i,j in enumerate(cols)]

[[1, 5], [7, 4], [3, 0]]

Answer 3

您可以将标签映射到索引，然后

take_along_axis

d = {c: i for i,c in enumerate(df.columns)}
idx = pd.DataFrame(cols).replace(d).to_numpy()
#array([[0, 1],
#       [2, 1],
#       [0, 1]])

np.take_along_axis(df.to_numpy(), idx, axis=1)
#array([[1, 5],
#       [7, 4],
#       [3, 0]])