Векторизованные способы заменить цикл при копировании столбцов в pandas

Как перенести значения между столбцами в pandas по словарю соответствий без вложенных циклов: два векторизованных подхода — DataFrame.assign и reindex+rename. Примеры и советы.

Когда нужно «нормализовать» DataFrame в pandas, перенося значения из «исходных» столбцов в несколько «целевых» по словарю соответствий, быстрый вложенный цикл решает задачу. Но если вам ближе идиоматичный, векторизованный стиль, хорошо сочетающийся с цепочками методов, у pandas есть подходящие инструменты.

Пример данных и базовая реализация

Рассмотрим DataFrame, где некоторые столбцы должны дублировать значения других в соответствии с простым отображением. Нужно переписать A2 и A3 значениями из A1, а A5 — значениями из A4.

import pandas as pd

frame = pd.DataFrame({
  "A1": [1, 11, 111],
  "A2": [2, 22, 222],
  "A3": [3, 33, 333],
  "A4": [4, 44, 444],
  "A5": [5, 55, 555]
})

rel_map = {
  "A1": ["A2", "A3"],
  "A4": ["A5"]
}

for src_col, targets in rel_map.items():
    for dst_col in targets:
        frame[dst_col] = frame[src_col]

print(frame)

Получаем ожидаемый результат:

    A1   A2   A3   A4   A5
0    1    1    1    4    4
1   11   11   11   44   44
2  111  111  111  444  444

В чем суть проблемы?

Вложенный цикл прозрачен и понятен, но выполняет работу на уровне итераций Python. В pandas операции, описывающие присваивание целиком, обычно короче, лучше встраиваются в рабочие цепочки и точнее соответствуют семантике DataFrame. Задача — заменить цикл векторизованным, «по‑пандасовому» выражением без изменения результата.

Векторизованные решения

Первый вариант использует DataFrame.assign, чтобы собрать все обновления целевых столбцов в одном вызове. Второй опирается на проекцию столбцов: reindex в паре с rename и set_axis выравнивает столбцы с помощью инвертированного отображения.

Используем assign с переписанным отображением:

updated_one = frame.assign(**{dest: frame.get(src)
                              for src, dest_list in rel_map.items()
                              for dest in dest_list})

print(updated_one)

Учтите, что assign не работает in place. Используйте его в цепочке или присвойте результат той же переменной, если хотите сохранить изменения.

Используем reindex и rename/set_axis, инвертировав отображение:

alias_map = {alias: origin
             for origin, aliases in rel_map.items()
             for alias in aliases}

updated_two = (
    frame.reindex(columns=frame.rename(columns=alias_map).columns)
         .set_axis(frame.columns, axis=1)
)

print(updated_two)

Оба подхода дают тот же результат, что и цикл:

    A1   A2   A3   A4   A5
0    1    1    1    4    4
1   11   11   11   44   44
2  111  111  111  444  444

Почему это важно

Когда трансформацию описываешь одним выражением, замысел становится очевидным, а решение легче вписывается в рабочие процессы pandas. Заодно сокращается объем итераций на стороне Python. Что до скорости, универсального ответа нет: относительная производительность зависит от общего числа столбцов, доли изменяемых столбцов, dtypes исходных данных и даже фрагментированности DataFrame. Надежнее всего — проверять на своих данных.

Выводы

Нужна наглядная, легко сочетаемая в цепочках конструкция — берите DataFrame.assign. Если удобнее управлять проекцией столбцов и важен точный контроль выравнивания, связка reindex плюс rename/set_axis тоже отлично работает. В обоих случаях целевые столбцы переписываются значениями из соответствующих источников — точно как в цикле. Когда важна скорость, замерьте оба варианта на реальной задаче и выберите тот, что лучше ведет себя в вашей среде.