How to Vectorise index insertion in Python?

Viewed 37

I have this problem with a very large data set which I can finished in Excel in less than a minute but it takes way too long on Python.

Objective: To give each row an ID based on information in column X and Y of the data set.

In Excel:

  • Initialize counter to 1
  • For each row i:
    • If both X = 0 and Y = 0, row ID = counter, followed by counter += 1
    • [edited] Else row ID = ID in the previous row.
  • Next i

My pd dataframe is large. Doing it in for loop takes more than an hour. I don't know how to vectorise my problem to avoid a for loop.

Hope someone can help me.

1 Answers

To find an efficient Pandas solution, you should rephrase your problem. Your counter is essentially the number of the previous all-zero rows (plus 1):

df = pd.DataFrame({'X': [0,2,1,0,0,1,2,0], 
                   'Y': [0,2,1,3,0,0,1,2]})
df['counter'] = (((df.X==0) & (df.Y==0)).cumsum().shift() + 1)\
                  .fillna(1).astype(int)
#   X  Y  counter
#0  0  0        1
#1  2  2        2
#2  1  1        2
#3  0  3        2
#4  0  0        2
#5  1  0        3
#6  2  1        3
#7  0  2        3
Related