How to count the distance in cells (e.g. in indices) between two repeating values in one column in Pandas dataframe?

Viewed 25

I have the following dataset. It lists the words that were presented to a participant in the psycholinguistic experiment (I set the order of the presentation of each word as an index):

data = {'Stimulus': ['sword','apple','tap','stick', 'elephant', 'boots', 'berry', 'apple', 'pear', 'apple', 'stick'],'Order': [1,2,3,4,5,6,7,8,9,10,11]}
df = pd.DataFrame(data, columns = ['Stimulus', 'Order'])
df.set_index('Order', inplace=True)


       Stimulus
Order   
1      sword
2      apple
3      tap
4      stick
5      elephant
6      boots
7      berry
8      apple
9      pear
10     apple
11     stick

Some values in this dataset are repeated (e.g. apple), some are not. The problem is that I need to calculate the distance in cells based on the order column between each occurrence of repeated values and store it in a separate column, like this:

       Stimulus  Distance
Order   
1      sword     NA
2      apple     NA
3      tap       NA
4      stick     NA
5      elephant  NA
6      boots     NA
7      berry     NA
8      apple     6
9      pear      NA
10     apple     2
11     stick     7

It shouldn't be hard to implement, but I've got stuck.. Initially, I made a dictionary of duplicates where I store items as keys and their indices as values:

{'apple': [2,8,10],'stick': [4, 11]}

And then I failed to find a solution to put those values into a column. If there is a simplier way to do it in a loop without using dictionaries, please let me know. I will appreciate any advice.

1 Answers

Use, df.groupby on Stimulus then transform the Order column using pd.Series.diff:

df = df.reset_index()
df['Distance'] = df.groupby('Stimulus').transform(pd.Series.diff)
df = df.set_index('Order')

# print(df)
       Stimulus  Distance
Order                    
1         sword       NaN
2         apple       NaN
3           tap       NaN
4         stick       NaN
5      elephant       NaN
6         boots       NaN
7         berry       NaN
8         apple       6.0
9          pear       NaN
10        apple       2.0
11        stick       7.0
Related