"Stretch" the data-frame and fill with zeros

user_11974225 0
Viewed 97

I have this data-frame:

ID  X  Var1  Var2  Var3
 A  1    52    16    17
 A  2    73     0    20
 A  3    60    42    16
 A  4    15    87    73
 A  5     0    18    63
 B  1    66    42     0
 B  3    13    28    64
 B  4     0    37     0
 C  1    22    16    18
 C  2    10    81    82
 C  3    11    84    48
 D  2    31    46    36
 D  4    74    31    49

I want that each ID will have rows with values ranging from 1 to 5 in column X. In the other rows, the values would be 0 if the corresponding value in X didn't exist.

Expected results:

ID  X  Var1  Var2  Var3
 A  1    52    16    17
 A  2    73     0    20
 A  3    60    42    16
 A  4    15    87    73
 A  5     0    18    63
 B  1    66    42     0
 B  2     0     0     0
 B  3    13    28    64
 B  4     0    37     0
 B  5     0     0     0
 C  1    22    16    18
 C  2    10    81    82
 C  3    11    84    48
 C  4     0     0     0
 C  5     0     0     0
 D  1     0     0     0
 D  2    31    46    36
 D  3     0     0     0
 D  4    74    31    49
 D  5     0     0     0
pandas numpy python
3 Answers

You could use the complete function from pyjanitor, which exposes explicitly missing rows of values:

 import janitor

 df.complete(columns=['ID', 'X']).fillna(0)

   ID  X  Var1  Var2  Var3
0   A  1  52.0  16.0  17.0
1   A  2  73.0   0.0  20.0
2   A  3  60.0  42.0  16.0
3   A  4  15.0  87.0  73.0
4   A  5   0.0  18.0  63.0
5   B  1  66.0  42.0   0.0
6   B  2   0.0   0.0   0.0
7   B  3  13.0  28.0  64.0
8   B  4   0.0  37.0   0.0
9   B  5   0.0   0.0   0.0
10  C  1  22.0  16.0  18.0
11  C  2  10.0  81.0  82.0
12  C  3  11.0  84.0  48.0
13  C  4   0.0   0.0   0.0
14  C  5   0.0   0.0   0.0
15  D  1   0.0   0.0   0.0
16  D  2  31.0  46.0  36.0
17  D  3   0.0   0.0   0.0
18  D  4  74.0  31.0  49.0
19  D  5   0.0   0.0   0.0
user_71757130

create an empty dfn with ID and 1 to 5 X, then join the raw df.

ID_list = df.ID.unique()
dfn = pd.DataFrame({'ID':ID_list.repeat(5),
                    'X': [1,2,3,4,5] * len(ID_list)})

# method1
df_result = df.join(dfn.set_index(['ID', 'X']), 
        on=['ID', 'X'], how='right').fillna(0)

# method2
dfn[['Var1', 'Var2', 'Var3']] = 0
df_result = pd.concat([df, dfn], ignore_index=True)
df_result = df_result.drop_duplicates(['ID', 'X']).sort_values(['ID', 'X'])

df_result

   ID  X  Var1  Var2  Var3
0   A  1    52    16    17
1   A  2    73     0    20
2   A  3    60    42    16
3   A  4    15    87    73
4   A  5     0    18    63
5   B  1    66    42     0
19  B  2     0     0     0
6   B  3    13    28    64
7   B  4     0    37     0
22  B  5     0     0     0
8   C  1    22    16    18
9   C  2    10    81    82
10  C  3    11    84    48
26  C  4     0     0     0
27  C  5     0     0     0
28  D  1     0     0     0
11  D  2    31    46    36
30  D  3     0     0     0
12  D  4    74    31    49
32  D  5     0     0     0
user_60063830

The most natural way to do your task is to group rows by ID and add "missing" rows to each group.

To do it, define the following function:

def addRows(grp):
    rng = range(1, 6)
    wrk = pd.DataFrame({'X': rng}, index=rng)
    rv = wrk.join(grp.set_index('X')).fillna(0, downcast='infer')
    rv.ID = grp.iloc[0,0]
    return rv

Then apply it to each group and reset the index:

result = df.groupby('ID', group_keys=False).apply(addRows).reset_index(drop=True)

The result is:

    X ID  Var1  Var2  Var3
0   1  A    52    16    17
1   2  A    73     0    20
2   3  A    60    42    16
3   4  A    15    87    73
4   5  A     0    18    63
5   1  B    66    42     0
6   2  B     0     0     0
7   3  B    13    28    64
8   4  B     0    37     0
9   5  B     0     0     0
10  1  C    22    16    18
11  2  C    10    81    82
12  3  C    11    84    48
13  4  C     0     0     0
14  5  C     0     0     0
15  1  D     0     0     0
16  2  D    31    46    36
17  3  D     0     0     0
18  4  D    74    31    49
19  5  D     0     0     0
user_73884770
Related