从pandas dataframe列中提取多个单词到同一列中

Question

Suppose a dataframe consists of two columns A={1,2,3} B={'abc d', 'efg h', 'ijk l'}. For A = 2, I would like to change the corresponding entry in column B to 'ef h'. (ie. extract the first, second and last word, not drop the third word, not the same).

It is easy to extract single words using the df.loc[df['colA']=2,'colB'].str.split().str[x], where x= 0,1 and -1, but I'm having difficulty joining the three words back into one string efficiently. The most efficient way I can think of is provided below. Is there a better way of achieving what I'm trying to do? Thanks.

y = lambda x : df.loc[df['colA']==2,'colB'].str.split().str[x]
df.loc[df['colA']=2,'colB'] = y(0) + ' ' + y(1) + ' ' + y(-1)

Expected and actual result:

A     B
1  a b c d
2  e f h
3  i j k l

Answer 1

How about this:

df = pd.DataFrame(data = {'A': [1,2,3], 
                          'B': ['a b c d', 'e f g h', 'i j k l']})

y = lambda x : df.loc[df['A']==2,'B'].str[0:2*x+2] + df.loc[df['A']==2,'B'].str[-1]
df.loc[df1['A']==2,'B'] = y(1)

Then df is the wanted:

   A        B
0  1  a b c d
1  2    e f h
2  3  i j k l

Answer 2

You were pretty close to the solution, the only problem is that str[x] returns a value wrapped in a Series object. You could fix this by extracting the value from the Series as shown:

y = lambda x : df.loc[df['colA']==2,'colB'].str.split().str[x].values[0]
df.loc[df['colA']==2,'colB'] = y(0) + ' ' + y(1) + ' ' + y(-1)

You can also achieve the same by making use of the apply function

df.loc[df['colA']==2, 'colB'] = df.loc[df['colA']==2,'colB'].apply(lambda x: ' '.join(x.split()[0:2] + [x.split()[-1]]))