Choose column 'row' value based on row value in panda data frame

Question

I got a data frame with multiple rows like this:

N  amount country codeA codeB
0  -15678   NAN    CH    US
1  3475     NAN    POL   FR

How can I select column codeA for the value of country when the amount is negative and codeB when positive

I would like to have something like this as output

N  amount country codeA codeB
0  -15678   CH    CH    US
1  3475     FR    POL   FR

Answer 1

Try np.where:

df['country'] = np.where(df['amount'] < 0, df['codeA'], df['codeB'])
print(df)

Prints:

   N  amount country codeA codeB
0  0  -15678      CH    CH    US
1  1    3475      FR   POL    FR

Answer 2

You can use the apply method along with a lambda function in pandas. :

import pandas as pd

# Create a sample DataFrame
data = {
    'N': [0, 1],
    'amount': [-15678, 3475],
    'country': ['NAN', 'NAN'],
    'codeA': ['CH', 'POL'],
    'codeB': ['US', 'FR']
}

df = pd.DataFrame(data)

# Apply lambda function to select codeA or codeB based on the amount value
df['country'] = df.apply(lambda row: row['codeA'] if row['amount'] < 0 else row['codeB'], axis=1)

print(df)

Output:

   N  amount country codeA codeB
0  0  -15678      CH    CH    US
1  1    3475      FR   POL    FR

Answer 3

You can create a custom function to calculate country value for each row and then apply it to the DataFrame row-by-row with .apply() method. axis=1 means that the fucntion is applied to each row (as opposed to each column).

import pandas as pd

df = pd.DataFrame({
    "N": [0, 1],
    "amount": [-15678, 3475],
    "codeA": ["CH", "POL"],
    "codeB": ["US", "FR"]
})

def get_country(row):
    return row["codeA"] if row["amount"] < 0 else row["codeB"]

df["country"] = df.apply(get_country, axis=1)
df

Output:

    N       amount  codeA   codeB   country
0   0       -15678  CH      US      CH
1   1       3475    POL     FR      FR

Answer 4

I think you can use this below code

df['country'] = df.apply(lambda x: x['codeA'] if x['amount'] < 0 else x['codeB'], axis=1)

Answer 5

You can achieve this by using the apply function along with a lambda function in pandas. try this below code:

import pandas as pd

# Create the DataFrame
data = {'N': [0, 1],
        'amount': [-15678, 3475],
        'country': ['NAN', 'NAN'],
        'codeA': ['CH', 'POL'],
        'codeB': ['US', 'FR']}

df = pd.DataFrame(data)

# Define a function to select the appropriate column based on the amount value
def select_column(row):
    if row['amount'] < 0:
        return row['codeA']
    else:
        return row['codeB']

# Apply the function to create a new column with the selected code
df['selected_code'] = df.apply(lambda row: select_column(row), axis=1)

# Update the 'country' column based on the selected code
df['country'] = df['selected_code']

# Drop the 'selected_code' column if not needed
df.drop('selected_code', axis=1, inplace=True)

# Print the resulting DataFrame
print(df)