How to select and combine different columns based on specific condition in pandas python?

Qendrim Krasniqi 0
Viewed 44
df = pd.DataFrame(data={
    "id": ['a', 'a', 'b', 'b', 'a', 'c', 'c', 'b'],
    "transaction_amount": [110, 0, 10, 30, 40.4, 62.2, 20, 20],
    "principal_amount":   [100, 0, 0,  0,  40,   60,   0,  0],
    "interest_amount":    [10,  0, 10, 0,  0.4,  0.6,  10, 0],
    "overpayment_amount": [0,   0, 0,  0,  0,    1.6,  10, 20],
})

I have the above dataframe. I want to have a column ,amount, and populate it as follows:

  • Create a row for each principal_amount, interest_amount and overpayment_amount if it's value is not 0, and assign principal, interest and overpayment to a new column, transaction_type, respectively.
  • Get value from transaction_amount if other three column values are 0 for that row.

The output should look like this:

   amount transaction_type id
3    30.0              NaN  b
0   100.0        principal  a
4    40.0        principal  a
5    60.0        principal  c
0    10.0         interest  a
2    10.0         interest  b
4     0.4         interest  a
5     0.6         interest  c
6    10.0         interest  c
5     1.6      overpayment  c
6    10.0      overpayment  c
7    20.0      overpayment  b

My current solution:

import pandas as pd

df = pd.DataFrame(data={
    "id": ['a', 'a', 'b', 'b', 'a', 'c', 'c', 'b'],
    "transaction_amount": [110, 0, 10, 30, 40.4, 62.2, 20, 20],
    "principal_amount":   [100, 0, 0,  0,  40,   60,   0,  0],
    "interest_amount":    [10,  0, 10, 0,  0.4,  0.6,  10, 0],
    "overpayment_amount": [0,   0, 0,  0,  0,    1.6,  10, 20],
})

columns = ["amount", "transaction_type"]
output_df = pd.DataFrame(columns=columns)

# Add transaction amount
condition = (df["principal_amount"] == 0) & (df["interest_amount"] == 0) & (df["overpayment_amount"] == 0) & (df["transaction_amount"] != 0)
subdf = df.loc[condition, ['id', 'transaction_amount']]
subdf = subdf.rename(columns={'transaction_amount': "amount"})
output_df = output_df.append(subdf)

# Add principal and interest
for field in ["principal_amount", "interest_amount", "overpayment_amount"]:
    subdf = df.loc[df[field] != 0, ['id', field]]
    subdf["transaction_type"] = field.split("_")[0]
    subdf = subdf.rename(columns={field: "amount"})
    output_df = output_df.append(subdf)

Is there any pandas feature that helps me do this implementation more concise and efficient?

pandas python
1 Answers

One approach could be as follows.

import pandas as pd
import numpy as np

out = df.reset_index(drop=False).melt(
    id_vars=['index'], 
    value_vars=list(df.columns)[1:], 
    var_name='transaction_type', 
    value_name='amount'
    ).set_index('index')

out = out[out['amount'].gt(0)]
out['v'] = out.index.value_counts()

out = out[out.v.eq(1) | 
          out.transaction_type.ne('transaction_amount')].drop('v', axis=1)

out['transaction_type'] = out['transaction_type']\
    .str.replace('_amount','').replace({'transaction':np.nan})

out = out.iloc[:,::-1]
out.index.name=None
out['id'] = df['id']

print(out)

   amount transaction_type id
3    30.0              NaN  b
0   100.0        principal  a
4    40.0        principal  a
5    60.0        principal  c
0    10.0         interest  a
2    10.0         interest  b
4     0.4         interest  a
5     0.6         interest  c
6    10.0         interest  c
5     1.6      overpayment  c
6    10.0      overpayment  c
7    20.0      overpayment  b

Explanation approach:

  • We use df.melt to get all the column names (starting from the second column) and amounts in two separate columns, and make sure also to keep the original index values (reset index first, and then set it to 'index' again).
  • We keep only the rows where amount > 0 by using Series.gt on amount.
  • We create a temporary column to store Series.value_counts applied to the index. Each index value with a value count of 1 will only have a value associated with transaction_amount.
  • We use this information for another filter: keep only rows that have out['v'].eq(1) or have a transaction_type that is not 'transaction_amount'. Afterwards, we can drop the temporary column again.
  • Finally, we get rid of '_amount' in the column transaction_type and also replace 'transaction' with NaN values. Last cosmetic procedure is to get the columns in the requested order, to remove the index name, and add id as an extra column.
ouroboros10
Related