Get sample index within dataasets' mapping function

Hello!

I’m trying to create a dictionary mapping classes to their corresponding samples’ index in the dataset, so that given a class I can quickly retrieved its samples. While I could simply do a for loop, I want to utilize the map func. Is there a way to get the index of the sample within the mapping function?

Thanks for helping!