pandas 基于另一列的值一次创建多个列

mum43rcc  于 2023-04-28  发布在  其他
关注(0)|答案(5)|浏览(145)

我有以下dataframe:

import pandas as pd
import random

# Create the lists for each column
nrows = 5
a = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
b = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
c = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
idx = [random.randint(0, 3) for i in range(nrows)]

# Create the pandas dataframe
df = pd.DataFrame({'a': a, 'b': b, 'c': c, 'idx': idx})

我想再创建3列a_desb_desc_des,方法是为每一行提取对应于该行idx值的abc值。(代码重复),需要更新的列越多,需要重复的代码就越多。是否可以用一条apply语句生成所有三列a_desb_desc_des

EDIT:抱歉,我犯了一个错误,列表的长度不一样,我今天没有时间来解决这个问题,但我明天一定会解决的。

tcomlyy6

tcomlyy61#

IIUC,你可以使用assigndict/listcomp & * unpacking *:

out = (df.assign(**{f"{col}_des": [e[i] for (e,i) in zip(df[col], df["idx"])]
                    for col in df.columns[:-1]})) #or `for col in ["a", "b", "c"]`

输出:

a           b            c  idx  a_des  b_des  c_des
0   [1, 9, 7]  [10, 7, 9]    [9, 9, 4]    1      9      7      9
1  [2, 10, 3]   [5, 8, 4]    [7, 0, 0]    1     10      8      0
2   [2, 8, 7]   [0, 8, 7]    [2, 9, 5]    0      2      0      2
3   [1, 2, 8]   [5, 5, 6]  [7, 10, 10]    1      2      5     10
4  [10, 0, 9]   [2, 0, 7]    [6, 6, 7]    1      0      0      6
m528fe3b

m528fe3b2#

我会说

import random
import pandas as pd

a = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
b = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
c = [[random.randint(0, 10), random.randint(0, 10), random.randint(0, 10)] for i in range(nrows)]
idx = [random.randint(0, 2) for i in range(nrows)]

df = pd.DataFrame({'a': a, 'b': b, 'c': c, 'idx': idx})

# Solution:
cols = ["a", "b", "c"] 
idx_col = "idx"

out = df.apply(lambda row: row[cols].str[row[idx_col]].add_suffix("_des"), axis=1)

结果:

>>> df
            a          b          c  idx
0   [4, 1, 2]  [1, 2, 9]  [8, 1, 5]    2
1   [6, 6, 1]  [4, 2, 5]  [4, 7, 1]    0
2   [8, 1, 7]  [7, 1, 1]  [9, 4, 1]    1
3  [6, 10, 4]  [6, 3, 0]  [6, 3, 1]    0
4  [10, 3, 5]  [4, 0, 3]  [6, 3, 7]    2
>>> out
   a_des  b_des  c_des
0      2      9      5
1      6      4      4
2      1      1      4
3      6      6      6
4      5      3      7

同样值得注意的是,如果你有一个3D numpy数组,你可以用向量的方式来做这件事,你可以考虑。

3df52oht

3df52oht3#

您可以编写一个函数,返回idx对应的值,然后使用apply()

def ret_idx(x):
    x.a = x.a[x.idx]
    x.b = x.b[x.idx]
    x.c = x.c[x.idx]
    return x[['a', 'b', 'c']]

df[['a_des', 'b_des', 'c_des']] = df.apply(ret_idx, axis=1)

df

a           b          c  idx  a_des  b_des  c_des
0   [6, 5, 0]   [3, 5, 1]  [4, 9, 4]    2      0      1      4
1   [4, 2, 1]   [0, 4, 2]  [7, 0, 6]    2      1      2      6
2   [7, 7, 2]  [10, 6, 0]  [5, 1, 8]    2      2      0      8
3  [2, 5, 10]   [4, 2, 1]  [4, 1, 6]    0      2      4      4
4   [4, 9, 1]  [4, 5, 10]  [9, 9, 7]    1      9      5      9
oknwwptz

oknwwptz4#

如果所有列表的长度都相同,请使用3D numpy数组:

cols = ['a', 'b', 'c']

df.join(pd.DataFrame(np.array(df[cols].to_numpy().tolist())
                     [np.arange(len(df)), :, df['idx']],
                    columns=cols
                    ).add_suffix('_des')
       )

输出:

a            b           c  idx  a_des  b_des  c_des
0  [9, 2, 10]    [8, 7, 7]   [6, 4, 5]    2     10      7      5
1   [3, 4, 7]  [4, 10, 10]  [8, 10, 3]    1      4     10     10
2  [10, 8, 8]    [1, 4, 6]   [1, 0, 5]    2      8      6      5
3  [2, 10, 9]    [4, 8, 7]   [4, 3, 1]    2      9      7      1
4  [8, 5, 10]   [0, 9, 10]  [5, 1, 10]    0      8      0      5
w1e3prcc

w1e3prcc5#

对于不同长度的列表:

def extract_index(df):
    col = df.columns[0]
    return (df.explode(col).assign(g=lambda x: x.groupby(level=0).cumcount())
              .query('idx == g')[col].rename(f'{col}_des'))

out = pd.concat([extract_index(df[[col, 'idx']])
                for col in df.drop(columns='idx')], axis=1)

输出:

>>> out
  a_des b_des c_des
0     4   NaN     6
1     2     5   NaN
2     3     7     0
3     1     7     4
4   NaN   NaN     8

>>> df
              a             b             c  idx
0  [3, 2, 4, 8]           [3]  [2, 1, 6, 3]    2
1  [1, 7, 2, 2]  [4, 0, 3, 5]        [0, 3]    3
2     [0, 3, 2]        [7, 7]  [7, 0, 3, 5]    1
3        [4, 1]        [9, 7]     [0, 4, 6]    1
4        [3, 8]           [1]  [9, 2, 9, 8]    3

IIUC,可以使用explode

>>> (df.explode(['a', 'b', 'c']).assign(g=lambda x: x.groupby(level=0).cumcount())
       .query('idx == g').drop(columns=['idx', 'g']).add_suffix('_des'))

  a_des b_des c_des
0     9     0     6
1     6     6     0
3     7     4     9
4     1     3     5

一步一步:

>>> out = df.explode(['a', 'b', 'c'])
   a   b   c  idx
0  9   4   2    1
0  9   0   6    1
0  5   5  10    1
1  6   5   1    1
1  6   6   0    1
1  6  10   4    1
2  9   1   9    3
2  7   2   3    3
2  2   3   9    3
3  1   9   7    1
3  7   4   9    1
3  4   6   8    1
4  9   4   7    1
4  1   3   5    1
4  9   5  10    1

>>> out = out.assign(g=lambda x: x.groupby(level=0).cumcount())
   a   b   c  idx  g
0  9   4   2    1  0
0  9   0   6    1  1
0  5   5  10    1  2
1  6   5   1    1  0
1  6   6   0    1  1
1  6  10   4    1  2
2  9   1   9    3  0
2  7   2   3    3  1
2  2   3   9    3  2
3  1   9   7    1  0
3  7   4   9    1  1
3  4   6   8    1  2
4  9   4   7    1  0
4  1   3   5    1  1
4  9   5  10    1  2

>>> out = out.query('idx == g').drop(columns=['idx', 'g']).add_suffix('_des')
  a_des b_des c_des
0     9     0     6
1     6     6     0
3     7     4     9
4     1     3     5

您可以使用pd.concat合并dfout

df = pd.concat([df, out], axis=1)

相关问题