Fraz Fraz - 3 months ago 12
Python Question

Creating more rows based on condition in pandas

I have a dataframe like following:

id, index, val_1, val_2
1, 1, 0.2, 0
1, 2, 0.4, 0.2
2,2, 0.1, 0.5
2,4, 0.7, 0.0
....


and so on

Now, the complete range of index values that are allowed for each id is

[1,2,3,4]


So, if any of this index is missing for each id, I want to add those rows.
So for the above example, the desired output is

id, index, val_1, val_2
1, 1, 0.2, 0
1, 2, 0.4, 0.2
1, 3, 0, 0 # added because index 3 was missing for id 1
1, 4, 0, 0 # added because index 4 was missing for id 1
2, 1,0,0 # added because index 1 was missing for id 2
2,2, 0.1, 0.5
2, 3, 0, 0
2,4, 0.7, 0.0
....


How do i perform this in operation in pandas?

Answer

try this:

In [210]: from itertools import product

In [211]: x = pd.DataFrame(list(product(df.id.unique(), [1,2,3,4])), columns=['id','index']).assign(val_1=0, val_2=0).set_index(['id','index'])

In [212]: x.update(df.set_index(['id','index']))

In [213]: x
Out[213]:
          val_1  val_2
id index
1  1        0.2    0.0
   2        0.4    0.2
   3        0.0    0.0
   4        0.0    0.0
2  1        0.0    0.0
   2        0.1    0.5
   3        0.0    0.0
   4        0.7    0.0

In [214]: x.reset_index()
Out[214]:
   id  index  val_1  val_2
0   1      1    0.2    0.0
1   1      2    0.4    0.2
2   1      3    0.0    0.0
3   1      4    0.0    0.0
4   2      1    0.0    0.0
5   2      2    0.1    0.5
6   2      3    0.0    0.0
7   2      4    0.7    0.0

Explanation:

In [225]: x = (pd.DataFrame(list(product(df.id.unique(), [1,2,3,4])), columns=['id','index'])
   .....:        .assign(val_1=0, val_2=0)
   .....:        .set_index(['id','index']))

In [226]: x
Out[226]:
          val_1  val_2
id index
1  1          0      0
   2          0      0
   3          0      0
   4          0      0
2  1          0      0
   2          0      0
   3          0      0
   4          0      0

In [227]: x.update(df.set_index(['id','index']))

In [228]: x
Out[228]:
          val_1  val_2
id index
1  1        0.2    0.0
   2        0.4    0.2
   3        0.0    0.0
   4        0.0    0.0
2  1        0.0    0.0
   2        0.1    0.5
   3        0.0    0.0
   4        0.7    0.0

In [229]: x.reset_index()
Out[229]:
   id  index  val_1  val_2
0   1      1    0.2    0.0
1   1      2    0.4    0.2
2   1      3    0.0    0.0
3   1      4    0.0    0.0
4   2      1    0.0    0.0
5   2      2    0.1    0.5
6   2      3    0.0    0.0
7   2      4    0.7    0.0
Comments