How to create a new column conditionally?

I want to add a new column called id to my pandas DataFrame. If the value of id is just a fixed number, then I can create a new column as follows:

id_value = 1
df.insert(0, 'id', id_value)

However, in my case I need to create batches of 3 rows, as follows:

id
1
1
1
2
2
2
3
3
3
...

How can I do it?

Answer

You can create array by np.arange with integer division by 3, so starting by 0, if need custom first number add id_value:

id_value = 1
df.insert(0, 'id', np.arange(len(df)) // 3 + id_value)

If there is default index values also working:

df = pd.DataFrame({'a':np.random.randint(10, size=10)})

id_value = 1
df.insert(0, 'id', df.index // 3 + id_value)
print (df)
   id  a
0   1  2
1   1  8
2   1  5
3   2  3
4   2  1
5   2  1
6   3  9
7   3  2
8   3  7
9   4  5