从pd.DataFrame中获取列标签的索引的高效方法

Question

我有一个方法，根据标签从 pandas dataframe 中获取列，但是使用 numpy 进行索引更快。 在 pandas 或 numpy 中是否有一种方式可以在不迭代的情况下从列标签转换为列索引？

DF_var = pd.DataFrame(np.random.random((5,10)), columns=["attr_%d" % _ for _ in range(10)])
query_cols = ["attr_2","attr_5","attr_6","attr_0"]
want_idx = [0,2,5,6]

# Something like np.where w/o iterating through? 
# np.where(query_cols in DF_var.columns)
# TypeError: unhashable type: 'list'

# np.where(x in DF_var.columns for x in query_cols)
# (array([0]),)


long_way = list()
for i, label in enumerate(DF_var.columns):
    if label in query_cols:
        long_way.append(i)
# print(sorted(long_way))
# [0, 2, 5, 6]

- O.rka

---

请参见：https://dev59.com/sWcs5IYBdhLWcg3wPxce - albert

这是针对单个值还是列表的？ - O.rka

@O.rka 单个，但你可以使用列表推导式来获取所有的索引。 - Alex

这个任务使用numpy完成，代码如下：np.argwhere(DF_var.columns.isin(["attr_2","attr_5","attr_6","attr_0"])).flatten()。 - sirfz

@O.rka 您可以使用 searchsorted 方法。新的答案已发布到链接的重复目标中。 - Divakar

Answer 1

short_way = [df.columns.get_loc(col) for col in query_cols]
print(sorted(short_way))
# outputs [0, 2, 5, 6]

- Alex

---

哦，好的，所以它仍然需要迭代吗？ - O.rka

@O.rka 是的。据我所知，没有向量化的方法。 - Alex