Получить значение первого столбца в Pandas DataFrame, где строка соответствует условию

Question

Скажем, у меня есть датафрейм панд, который выглядит следующим образом:

   color  number
0    red       3
1   blue       4
2  green       2
3   blue       2

Я хочу получить первое значение из числового столбца, в котором цветовой столбец имеет значение 'blue', которое в этом случае вернет 4.

Я знаю, что это можно сделать с помощью loc примерно так:

df[df['color'] == 'blue']['number'][0]

Мне интересно, есть ли более оптимальный подход, учитывая, что мне нужно только первое появление.

Answer 1

Использовать head - это вернет первую строку, если цвет существует, и пустую Series в противном случае.

col = 'blue'
df.query('color == @col').head(1).loc[:, 'number']

1    4
Name: number, dtype: int64

В качестве альтернативы, чтобы получить один предмет, используйте obj.is_empty:

u = df.query('color == @col').head(1)
if not u.is_empty:
    print(u.at[u.index[0], 'number'])

# 4

---

Разница между head и idxmax для недопустимого цвета:

df.query('color == "blabla"').head(1).loc[:, 'number']
# Series([], Name: number, dtype: int64)

df.loc[(df['color'] == 'blabla').idxmax(),'number']
# 3

Answer 2

Используя iloc с np.where:

idx = next(iter(df['number'].iloc[np.where(df['color'].eq('blue'))]), -1)  # 4

Обратите внимание, что это также обрабатывает случай, когда цвет не существует. Для сравнения, df['color'].eq('orange').idxmax() дает 0, хотя 'orange' не существует в серии. Вышеуказанная логика даст -1.

numba

Мне интересно, есть ли более оптимальный подход, учитывая, что я только когда-нибудь нужно первое вхождение.

Да! Для более эффективного решения см. Эффективно вернуть индекс первого значения, удовлетворяющего условию в массиве . Numba позволяет эффективно выполнять итерацию по строкам. В этом случае вам нужно будет сначала разложить ваши строки , чтобы вы передавали числовые массивы только в Numba:

from numba import njit

# factorize series, pd.factorize maintains order,
# i.e. first item in values gives 0 index
idx, values = pd.factorize(df['color'])
idx_search = np.where(values == 'blue')[0][0]

@njit
def get_first_index_nb(A, k):
    for i in range(len(A)):
        if A[i] == k:
            return i
    return -1

res = df['number'].iat[get_first_index_nb(idx, 1)]  # 4

Конечно, для одноразового расчета это неэффективно. Но для последовательных вычислений решение, скорее всего, будет фактором быстрее, чем решения, которые проверяют равенство по всей серии / массиву.

Answer 3

Использование idxmax

df.loc[(df['color'] == 'blue').idxmax(),'number']
Out[698]: 4