Merge and align columns based on 2 columns

Viewed 50

I have a df like below:

ID  P1   P2   P3_1_A  P3_2_B   P4_1_A  P4_2_B
1   110  111   1        0       1       1
2   111  112   0        0       1       0
3   110  112   0        1       0       0
4   112  111   1        1       1       1

So, P1 and P2 has 3 distinct values 110, 111 and 112.

Column P3's and P4's are always equal in number (2 in the above case) but may vary from data to data.

I want a resultant data frame where all the P3 values are aligned to P1 and all P4 values to P2 and columns renamed and looks like below

ID   P     P_A   P_B
1    110    1     0
1    111    1     1
2    111    0     0
2    112    1     0
3    110    0     1
3    112    0     0
4    112    1     1
4    111    1     1

I know how to merge P1 and P2 and get to column P but not getting how to align P3 and P4 to P1 and P2 and get to P_1 and P_2

2 Answers

You can rename columns for _ substrings, so possible convert ID to index and split all values to MultiIndex and last reshape by DataFrame.stack:

df1 = df.rename(columns={'P1':'P3_','P2':'P4_'}).set_index('ID')
df1.columns = df1.columns.str.split('_', expand=True, n=1)
df1 = df1.stack(0).add_prefix('P').reset_index(level=1, drop=True).reset_index()
print (df1)
   ID    P  P1_A  P2_B
0   1  110     1     0
1   1  111     1     1
2   2  111     0     0
3   2  112     1     0
4   3  110     0     1
5   3  112     0     0
6   4  112     1     1
7   4  111     1     1

EDIT: For more general solution is possible extract columns names without double _ and pass to set_index:

print (df)
   ID   P1   P2  P3_1_A  P3_2_B  P4_1_A  P4_2_B  A
0   1  110  111       1       0       1       1  9
1   2  111  112       0       0       1       0  7
2   3  110  112       0       1       0       0  8
3   4  112  111       1       1       1       1  7

df1 = df.rename(columns={'P1':'P3__','P2':'P4__'})

cols = df1.columns[df1.columns.str.count('_') != 2]
df1 = df1.set_index(cols.tolist())
df1.columns = df1.columns.str.split('_', expand=True, n=1)
df1 = df1.stack(0).add_prefix('P').reset_index(level=-1, drop=True).reset_index()
print (df1)
   A  ID  P1_A  P2_B   P_
0  9   1     1     0  110
1  9   1     1     1  111
2  7   2     0     0  111
3  7   2     1     0  112
4  8   3     0     1  110
5  8   3     0     0  112
6  7   4     1     1  112
7  7   4     1     1  111

You can use the pivot_longer function from pyjanitor; at the moment you have to install the latest development version from github :

The columns have patterns (some end with a number (1 or 2), while others end with either 'A' or 'B'). We can take advantage of this pattern and reshape the data. We pass a list of new column names to names_to, and pass a list of regexes that match the patterns to names_pattern:

 # install the latest dev version of pyjanitor
 # pip install git+https://github.com/ericmjl/pyjanitor.git

 import janitor

df.pivot_longer(
    index="ID",
    names_to=("P", "P_A", "P_B"),
    names_pattern=("^P\d$", ".*A$", ".*B$"),
    sort_by_appearance=True,
)

   ID   P   P_A P_B
0   1   110 1   0
1   1   111 1   1
2   2   111 0   0
3   2   112 1   0
4   3   110 0   1
5   3   112 0   0
6   4   112 1   1
7   4   111 1   1

In the code above, the new P column takes all values from the previous dataframe that had columns P1 or P2(they end with numbers), P_A takes columns that end with 'A', while P_B takes columns that end with 'B'.

Related