Получить индекс из подмножества pandas multindex
Учитывая это pandas.core.frame.DataFrame
образец:
col
0 2 0
0 1
3 1
Я хотел бы получить те индексы, где col равен 1:
df[df['col']==1]
Ожидаемый результат будет вектор (0,3)
,
2 ответа
Решение
Вы можете использовать NumPy для создания массива индексов, а затем отфильтровать второй индекс:
import numpy as np
df = pd.DataFrame({'col': [0, 1, 1]},
index=pd.MultiIndex.from_tuples([(0, 2), (0, 0), (0, 3)]))
# col
# 0 2 0
# 0 1
# 3 1
idx = np.array(df[df['col'] == 1].index.tolist())
# array([[0, 0],
# [0, 3]])
res = idx[:, 1]
# array([0, 3])
Это может сделать трюк:
import pandas as pd
d = {'base':[2,0,3],'col':[0,1,1]}
df = pd.DataFrame(data=d)
col_is_1 = df['col'] == 1
df_temp = df[col_is_1]
matrix = df_temp.values
vec = matrix[:,0]
print(vec)
возврат [0,3] как и ожидалось