Create a pandas column of numbers from 1 to 3 and repeat again

user_14088584 0
Viewed 1066

I have a dataframe:

   StoreNumber    Year  
    1000          2000  
    1000          2001  
    1000          2002  
    1001          2000  
    1001          2001  
    1001          2002

I want to add a column so that my final dataframe looks like:

StoreNumber       Year   New
    1000          2000    1
    1000          2001    2
    1000          2002    3
    1001          2000    1
    1001          2001    2
    1001          2002    3

I don't want the new row to depend upon StoreNumber which looks ovious in example. I want to start numbering with 1 and when I reach 3, then again start with 1. How do I do it?

pandas dataframe python
4 Answers

You can use itertools.cycle to create an iterator and use it to generate the target sequence:

from itertools import cycle

num_cycle = cycle([1, 2, 3])
df['New'] = [next(num_cycle) for num in range(len(df))]

import pandas as pd 
import itertools 

df = pd.DataFrame(
    data = [
        (1000, 2000),
        (1000, 2001),
        (1000, 2002),
        (1001, 2000),
        (1001, 2001),
        (1001, 2002),
    ],
    columns=['StoreNumber', 'Year']
)

num_cycle = itertools.cycle([1, 2, 3])
df['New'] = [next(num_cycle) for num in range(len(df))]

print(df)

And the output will be

   StoreNumber   Year  New
0         1000   2000    1
1         1000   2001    2
2         1000   2002    3
3         1001   2000    1
4         1001   2001    2
5         1001   2002    3
user_71317570

You could make the base list [1, 2, 3] and repeat it as many times as necessary.

baselist = [1, 2, 3]
size = df.size[0]
df['New'] = (baselist * (size // len(baselist) + 1))[size]
user_66928980

You can use numpy.tile:

In [507]: import numpy as np

In [508]: list_int = [1,2,3]
In [510]: df['New'] = np.tile(list_int, len(df)//len(list_int) + 1)[:len(df)]

In [511]: df
Out[511]: 
   StoreNumber  Year  New
0         1000  2000    1
1         1000  2001    2
2         1000  2002    3
3         1001  2000    1
4         1001  2001    2
5         1001  2002    3
user_58208140

You can use np.r_ to generate the range then take modulo division by 3 and add 1 to create a counter which repeats every three rows:

df['New'] = np.r_[:len(df)] % 3 + 1

   StoreNumber  Year  New
0         1000  2000    1
1         1000  2001    2
2         1000  2002    3
3         1001  2000    1
4         1001  2001    2
5         1001  2002    3
user_128331660
Related