Python: выбор строк по значению в большом фрейме данных

Question

С учетом фрейма данных df:

 Column A: [0, 1, 3, 4, 6]

 Column B: [0, 0, 0, 0, 0]

Цель состоит в том, чтобы условно заменить значения в столбце B. Если значения столбца A существуют в наборе assginedToA, мы заменим соответствующие значения в столбце B напостоянная b.

Например: если b = 1 и assignToA = {1,4}, результат будет

Column A: [0, 1, 3, 4, 6]

Column B: [0, 1, 0, 1, 0]

---

Мой код для поиска значений A и записи значений B в него выглядит следующим образом:

df.loc[df['A'].isin(assignedToA),'B']=b

Этот код работает, но это действительно медленно для огромного информационного кадра.Есть ли у вас какие-либо советы, как ускорить этот процесс?

---

Фрейм данных имеет около 5 миллионов строк, а assignedToA имеет максимум 7 значений.

Answer 1

Вы можете найти улучшение производительности, опустившись до numpy:

df = pd.DataFrame({'A': [0, 1, 3, 4, 6],
                   'B': [0, 0, 0, 0, 0]})

def jp(df, vals, k):
    B = df['B'].values
    B[np.in1d(df['A'], list(vals))] = k
    df['B'] = B
    return df

def original(df, vals, k):
    df.loc[df['A'].isin(vals),'B'] = k
    return df

df = pd.concat([df]*100000)

%timeit jp(df, {1, 4}, 1)        # 8.55ms
%timeit original(df, {1, 4}, 1)  # 16.6ms