Loop to create a new row for a given range of numbers

user_6567270 0
Viewed 582

I am sure this is an easy fix, but I haven't been able to find the exact solution to my problem. My data set has a column called 'LANE' which contains unique values. I want to add rows for each 'LANE' based on a range of numbers (which would be 0 to 12). As a result each 'LANE' would have 13 rows with a new column 'NUMBER' ranging from 0 up to and including 12.

Example:

Input

LANE
a
b

Output

LANE    NUMBER
a       0
a       1
a       2
a       3
a       4
a       5
a       6
a       7
a       8
a       9
a       10
a       11
a       12
b       0
b       1
b       2
b       3
b       4
b       5
b       6
b       7
b       8
b       9
b       10
b       11
b       12

I am currently trying different forms of:

num = 0

while num <= 12:
    for x in df['LANE']:
        df['NUMBER'] = num
    num += 1

The problem with this loop is, I still have one record for each lane and the 'NUMBER' column only has the value 12.

pandas python loops
2 Answers

Comprehension

For loops are the natural and naive way to produce Cartesian products. Comprehensions allow us to embed this more succinctly.

pd.DataFrame(
    [[l, n] for l in df.LANE for n in range(12)],
    columns=['LANE', 'NUMBER']
)

   LANE  NUMBER
0     a       0
1     a       1
2     a       2
3     a       3
4     a       4
5     a       5
6     a       6
7     a       7
8     a       8
9     a       9
10    a      10
11    a      11
12    b       0
13    b       1
14    b       2
15    b       3
16    b       4
17    b       5
18    b       6
19    b       7
20    b       8
21    b       9
22    b      10
23    b      11

itertools.product

This logic is almost identical to the Comprehension solution but it uses itertools built in product function. product is an iterator that pops out each combination one at a time. I force the result by unpacking with the splat * like so [*product(a, b)]. Ultimately, it is a list of lists that gets passed to the pd.DataFrame constructor in the same way as the Comprehension solution above.

from itertools import product

pd.DataFrame([*product(df.LANE, range(12))], columns=['LANE', 'NUMBER'])

groupby/cumcount and repeat

I don't like this answer but it provides some perspective on the simplicity of the other answers.

I use repeat to replicate each index value 12 times. I use this repeated index in a loc which returns a dataframe sliced with passed index. I then use groupbys cumcount to count each position within the group and add that as a new column.

df.loc[df.index.repeat(12)].assign(NUMBER=lambda d: d.groupby('LANE').cumcount())

  LANE  NUMBER
0    a       0
0    a       1
0    a       2
0    a       3
0    a       4
0    a       5
0    a       6
0    a       7
0    a       8
0    a       9
0    a      10
0    a      11
1    b       0
1    b       1
1    b       2
1    b       3
1    b       4
1    b       5
1    b       6
1    b       7
1    b       8
1    b       9
1    b      10
1    b      11
user_23366540

Another approach using pandas as below:

# First approach, one liner code
df = pd.DataFrame({'Lane': ['a'] * 12 + ['b'] * 12,
                   'Number': list(range(12)) * 2})

# Second approach
df = pd.DataFrame({'Lane': ['a'] * 12 + ['b'] * 12})
df['Number'] = df.groupby('Lane').cumcount()
user_92303740
Related