У меня есть dataframe, как это:
col1 x y z
A yes no yes
B no no yes
C no yes no
D yes no yes
E no no yes
F yes yes no
Я хотел бы выбрать данные, подобные этому. Если мои критерии - найти все yes
для A
, я должен получить [x,z]
, то есть значения для A
которые когда-либо были yes
Если B
, [z]
C
должно дать [y]
Что делать?
Сначала создайте индекс по столбцу col1
для индексирования по loc
, сравните по значению и последним значениям index
для list
:
df = df.set_index('col1')
def get_val(df, idx, val):
a = df.loc[idx].eq(val)
return a.index[a].tolist()
print (get_val(df, 'A', 'yes'))
['x', 'z']
print (get_val(df, 'B', 'yes'))
['z']
print (get_val(df, 'C', 'yes'))
['y']
Вот еще один способ создания функции:
df.set_index('col1', inplace=True)
def find_yes(df, x):
return df.columns[df.loc[x] == 'yes'].tolist()
полный пример
import pandas as pd
data = '''\
col1 x y z
A yes no yes
B no no yes
C no yes no
D yes no yes
E no no yes
F yes yes no'''
fileobj = pd.compat.StringIO(data)
df = pd.read_csv(fileobj, sep='\s+')
df.set_index('col1', inplace=True)
def find_yes(df, x):
return df.columns[df.loc[x] == 'yes'].tolist()
print(find_yes(df, 'A'))
print(find_yes(df, 'B'))
Возвращает:
['x', 'z']
['z']
Вы можете использовать
In [499]: df.eq('yes').dot(df.columns)[df.col1.eq('A')]
Out[499]:
0 xz
dtype: object
In [500]: df.eq('yes').dot(df.columns)[df.col1.eq('B')]
Out[500]:
1 z
dtype: object
In [501]: df.eq('yes').dot(df.columns)[df.col1.eq('C')]
Out[501]:
2 y
dtype: object
Еще один вариант для вас - как насчет использования melt
затем groupby
:
from io import StringIO
import pandas as pd
data = StringIO('''col1 x y z
A yes no yes
B no no yes
C no yes no
D yes no yes
E no no yes
F yes yes no''')
df = pd.read_csv(data, sep='\s+')
m = df.melt(id_vars='col1')
matches = m[m['value'] == 'yes'].groupby('col1')\
.agg({'variable': list})
это дает следующий фрейм данных:
variable
col1
A [x, z]
B [z]
C [y]
D [x, z]
E [z]
F [x, y]
df =df[df.index.str.contain(val)]
a = df.eq(val).any()
print (a.index[a].tolist())