chouchou
chouchou

Reputation: 25

pandas: Getting the value of next column, after finding string in multiple columns

How does one get next column, after a string search in multiple columns?

My data has a various length of sets as below. I want to find 'AA' in 'n(index)' column and get the value in 'v(index)' which is just next to it.

df = pd.DataFrame(columns = ['n1', 'v1', 'n2', 'v2', 'n3', 'v3', 'n4', 'v4'])
df.loc[0]=['BB', 22, 'AA', 80,'BA', 20, 'AG', 50]
df.loc[1]=['AV', 90, 'AA', 2, np.nan, np.nan, np.nan, np.nan]
df.loc[2]=['AA', 10, 'DD', 9, 'PP', 12, np.nan, np.nan]
df.loc[3]=['AA', 50, 'AB',30, 'BV',30, np.nan, np.nan]
print(df)

n1  v1  n2  v2  n3  v3  n4  v4
0   BB  22  AA  80  BA  20  AG  50
1   AV  90  AA  2   NaN NaN NaN NaN
2   AA  10  DD  9   PP  12  NaN NaN
3   AA  50  AB  30  BV  30  NaN NaN

I tried

df['AA'] = (df.values == 'AA').shift(1, axis=1).astype(int)

which does not work. How can I make the data like below?

    n1  v1  n2  v2  n3  v3  n4  v4  AA
0   BB  22  AA  80  BA  20  AG  50  80
1   AV  90  AA  2   NaN NaN NaN NaN 2
2   AA  10  DD  9   PP  12  NaN NaN 10
3   AA  50  AB  30  BV  30  NaN NaN 50

Upvotes: 1

Views: 522

Answers (1)

sammywemmy
sammywemmy

Reputation: 28709

Search for the positions in the dataframe for 'AA':

location = np.argwhere(df.isin(["AA"]).to_numpy())
location
array([[0, 2],
       [1, 2],
       [2, 0],
       [3, 0]])

Next, add 1 to the column values in the location array, since you are interested in the adjacent values:

location[:, -1] = location[:, -1] + 1
location
array([[0, 3],
       [1, 3],
       [2, 1],
       [3, 1]])

Get your values:

adjacent_values = [df.iat[x, y] for x, y in location]
adjacent_values
[80, 2, 10, 50]

Assign to the column:

df.assign(AA = adjacent_values)
    n1  v1  n2  v2  n3  v3  n4  v4  AA
0   BB  22  AA  80  BA  20  AG  50  80
1   AV  90  AA  2   NaN NaN NaN NaN 2
2   AA  10  DD  9   PP  12  NaN NaN 10
3   AA  50  AB  30  BV  30  NaN NaN 50

Upvotes: 1

Related Questions