Python/Pandas：使用一個列的值作為我想要一個值的列名的后綴

Question

我有一個 pandas dataframe。 從其中的多列中，我需要根據該行的 ID（在本例中為bar ），將 select 的值從僅一列轉換為單個新列。 我需要最快的方法來做到這一點。

Dataframe 的應用是這樣的：

foo bar ID_A    ID_B    ID_C    ID_D    ID_E    ...
1   B   1.5     2.3     4.1     0.5     6.6     ...
2   E   3       4       5       6       7       ...
3   A   9       6       3       8       1       ...
4   C   13      5       88      9       0       ...
5   B   6       4       6       9       4       ...
...

因此，一種方法的示例（我目前最快的）是 - 但是，對於我的目的來說它太慢了。

df.loc[df.bar=='A', 'baz'] = df.ID_A
df.loc[df.bar=='B', 'baz'] = df.ID_B
df.loc[df.bar=='C', 'baz'] = df.ID_C
df.loc[df.bar=='D', 'baz'] = df.ID_D
df.loc[df.bar=='E', 'baz'] = df.ID_E
df.loc[df.bar=='F', 'baz'] = df.ID_F
df.loc[df.bar=='G', 'baz'] = df.ID_G

結果將是這樣的（刪除使用的列之后）：

foo baz
1   2.3
2   7
3   9
4   88
5   4
...

我已經嘗試過.apply()並且速度很慢。
我嘗試使用np.where() ，它仍然比上面顯示的示例慢得多（比np.where()快 1000%）。

將不勝感激的建議！ 非常感謝

編輯：在前幾個答案之后，我想我需要添加這個：
“雖然我很欣賞相對於示例的運行時估計，但我知道這是一個小示例，因此可能會很棘手。我的實際數據有 280000 行和額外的 50 列（我需要與foo和baz保持一致）。我必須每個示例將 13 列減少到單列。速度是詢問的唯一原因，到目前為止，在最初的幾個回復中沒有提到速度。再次感謝！”

Answer 1

您可以使用索引查找的變體：

idx, cols = pd.factorize('ID_'+df['bar'])
out = pd.DataFrame({'foo': df['foo'],
                    'baz': df.reindex(cols, axis=1).to_numpy()[np.arange(len(df)), idx]})

output：

   foo   baz
0    1   2.3
1    2   7.0
2    3   9.0
3    4  88.0
4    5   4.0

測試速度

設置一個測試數據集（280k 行，54 個 ID 列）：

from string import ascii_uppercase, ascii_lowercase

letters = list(ascii_lowercase+ascii_uppercase)
N = 280_000

np.random.seed(0)

df = (pd.DataFrame({'foo': np.arange(1, N+1),
                    'bar': np.random.choice(letters, size=N)})
        .join(pd.DataFrame(np.random.random(size=(N, len(letters))),
                           columns=[f'ID_{l}' for l in letters]
                          ))
     )

速度測試：

%%timeit
idx, cols = pd.factorize('ID_'+df['bar'])
out = pd.DataFrame({'foo': df['foo'],
                    'baz': df.reindex(cols, axis=1).to_numpy()[np.arange(len(df)), idx]})

output：

54.4 ms ± 3.22 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 2

可以試試這個。 它應該推廣到任意數量的列。

import pandas as pd
import numpy as np

df = pd.DataFrame([[1, 'B', 1.5, 2.3, 4.1, 0.5, 6.6],
      [2, 'E', 3, 4, 5, 6, 7],
      [3, 'A', 9, 6, 3, 8, 1],
      [4, 'C', 13, 5, 88, 9, 0],
      [5, 'B', 6, 4, 6, 9, 4]])

df.columns = ['foo', 'bar', 'ID_A', 'ID_B', 'ID_C', 'ID_D', 'ID_E']

for val in np.unique(df['bar'].values):
    df.loc[df.bar==val, 'baz'] = df[f'ID_{val}']

Answer 3

為了展示另一種方法，您可以執行融合數據和重新索引的組合。 在這種情況下，由於列名的圖案性質，我使用wide_to_long （而不是熔化/堆棧）：

out = (
    pd.wide_to_long(
        df, stubnames=['ID'], i=['foo', 'bar'], j='', sep='_', suffix=r'\w+'
    )
    .loc[lambda d: 
        d.index.get_level_values('bar') == d.index.get_level_values(level=-1),
        'ID'
    ]
    .droplevel(-1)
    .rename('baz')
    .reset_index()
)

print(out)
   foo bar   baz
0    1   B   2.3
1    2   E   7.0
2    3   A   9.0
3    4   C  88.0
4    5   B   4.0

---

上述替代語法利用.melt和.query來縮短語法。

out = (
    df.melt(id_vars=['foo', 'bar'], var_name='id', value_name='baz')
    .assign(id=lambda d: d['id'].str.get(-1))
    .query('bar == id')
)

print(out)
    foo bar id   baz
2     3   A  A   9.0
5     1   B  B   2.3
9     5   B  B   4.0
13    4   C  C  88.0
21    2   E  E   7.0