Я пытаюсь создать в моем DataFrame новый столбец, представляющий собой список агрегированных имен столбцов. Вот пример DataFrame:

In [1]: df = pd.DataFrame({'A':[1,2,3],
               'B':[4,5,6],
               'C':[7,8,9],
               'D':[1,3,5],
               'E':[5,3,6],
               'F':[7,4,3]})
In [2]: df
Out[2]:
   A  B  C  D  E  F
0  1  4  7  1  5  7
1  2  5  8  3  3  4
2  3  6  9  5  6  3

Я хотел бы создать новый столбец, содержащий список имен столбцов, где выполняется определенное условие. Скажите, что меня интересуют столбцы, где значение> 3 - я бы хотел вывод, который выглядит следующим образом:

In [3]: df
Out[3]:
   A  B  C  D  E  F  Flag
0  1  4  7  1  5  7  ['B', 'C', 'E', 'F']
1  2  5  8  3  3  4  ['B', 'C', 'F']
2  3  6  9  5  6  3  ['B', 'C', 'D', 'E']

В настоящее время я использую apply:

df['Flag'] = df.apply(lambda row: [list(df)[i] for i, j in enumerate(row) if j > 3], axis = 1)

Это делает работу, но чувствует себя неуклюже, и мне интересно, есть ли более элегантное решение.

Благодарность!

2
Le Chase 1 Май 2019 в 18:26

7 ответов

Лучший ответ

Используйте df.dot() Вот:

df['Flag']=(df>3).dot(df.columns).apply(list)
print(df)

   A  B  C  D  E  F          Flag
0  1  4  7  1  5  7  [B, C, E, F]
1  2  5  8  3  3  4     [B, C, F]
2  3  6  9  5  6  3  [B, C, D, E]
4
anky_91 1 Май 2019 в 15:28

Изменить : добавление времени для всех решений этого вопроса

Я предпочитаю решение без apply

df['Flag'] = df.reset_index().melt(id_vars='index', value_name='val', var_name='col').query('val > 3').groupby('index')['col'].agg(list)

Или же

df['Flag'] = df.stack().rename('val').reset_index(level=1).query('val > 3').groupby(level=0)['level_1'].agg(list)

Out[2576]:
   A  B  C  D  E  F          Flag
0  1  4  7  1  5  7  [B, C, E, F]
1  2  5  8  3  3  4     [B, C, F]
2  3  6  9  5  6  3  [B, C, D, E]

Данные испытаний:

a = [
    [1,  4,  7,  1,  5,  7],
    [2,  5,  8,  3,  3,  4],
    [3,  6,  9,  5,  6,  3],
    ] * 10000

df = pd.DataFrame(a, columns = list('ABCDEF'))  

Время с %timeit:

In [79]: %timeit (df>3).dot(df.columns).apply(list)
40.8 ms ± 1.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [80]: %timeit [df.columns[x].tolist() for x in df.gt(3).values]
1.23 s ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [81]: %timeit df.gt(3).dot(df.columns).apply(list)
37.6 ms ± 644 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [82]: %timeit df.T.apply(lambda x: list(x[x>3].index))
16.4 s ± 99.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [83]: %timeit df.stack().rename('val').reset_index(level=1).query('val > 3')
    ...: .groupby(level=0)['level_1'].agg(list)
4.05 s ± 15.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [84]: %timeit df.apply(lambda x: df.columns[np.argwhere(x>3).ravel()].values
    ...: , 1)
c:\program files\python37\lib\site-packages\numpy\core\fromnumeric.py:56: Future
Warning: Series.nonzero() is deprecated and will be removed in a future version.
Use Series.to_numpy().nonzero() instead
  return getattr(obj, method)(*args, **kwds)
12 s ± 45.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Наиболее быстрым является решение с использованием .dot

1
Andy L. 1 Май 2019 в 18:26

Один из вариантов - создать кадр данных booleans, проверив, какие значения превышают определенный порог, с помощью DataFrame.gt и возьмите dot с именами столбцов. Наконец, используйте apply(list) для получения списков из результирующих строк:

df['Flag'] = df.gt(3).dot(df.columns).apply(list)

   A  B  C  D  E  F          Flag
0  1  4  7  1  5  7  [B, C, E, F]
1  2  5  8  3  3  4     [B, C, F]
2  3  6  9  5  6  3  [B, C, D, E]
2
yatu 1 Май 2019 в 15:36

Я все еще люблю здесь цикл

df['Flag']=[df.columns[x].tolist() for x in df.gt(3).values]
df
Out[968]: 
   A  B  C  D  E  F          Flag
0  1  4  7  1  5  7  [B, C, E, F]
1  2  5  8  3  3  4     [B, C, F]
2  3  6  9  5  6  3  [B, C, D, E]
4
YO and BEN_W 1 Май 2019 в 15:43

Мы также можем использовать @

df['Flag'] = ((df >3) @ df.columns).map(list)
0
Akhilesh 7 Май 2019 в 18:15

Использование numpy.argwhere и ravel() :

   df.apply(lambda x: df.columns[np.argwhere(x>3).ravel()].values, 1)
0
hacker315 1 Май 2019 в 18:08

По-другому:

df['Flag'] = df.T.apply(lambda x: list(x[x>3].index))
1
Quang Hoang 1 Май 2019 в 15:33