CODEWITHSUNDEEP
pythonpandasdataframe
martina.physics
martina.physics

Reputation: 9804

Pandas: grouping on couples from same column

Suppose I have a table (a dataframe) which contains this data:

| user     | food          |
|:--------:|:-------------:|
| 'A'      | 'meat'        | 
| 'A'      | 'carrot'      |
| 'A'      | 'candy'       |
| 'B'      |  'meat'       |
| 'B'      |  'carrot'      |
| 'C'      |  'meat'       |
| 'C'      |  'carrot'     |

Code:

df = pd.DataFrame({
    "user":["A", "A", "A", "B", "B", "C", "C"],
    "food":['meat', 'carrot', 'candy', 'meat', 'carrot', 'meat', 'carrot']
})

What I'd like to build is a table which for each couple of foods tells me the number of users which have them:

| food 1     | food 2        |  num users | 
|:----------:|:-------------:|:----------:| 
| 'meat'     | 'carrot'      | 3          | 
| 'meat'     | 'candy'       | 1          | 
| 'carrot'     | 'candy'       | 1          |

Is there a way to do this?

Upvotes: 2

Views: 65

Answers (2)

eugenhu
eugenhu

Reputation: 1238

You can try this:

food_pairs = [("meat", "carrot"), ("meat", "candy")]

food_to_users = {food: set(df.user[df.food == food].unique()) for food in df.food.unique()}

out = pd.DataFrame(
    ((*pair, len(set.intersection(*(food_to_users[food] for food in pair)))) for pair in food_pairs),
    columns=["food1", "food2", "num users"]
)

Average run time over 1000 trials was 0.00256s.

Test code for scalability:

import itertools
import math
import pandas as pd
from random import shuffle
from timeit import time

SIZE_OF_PAIRS = 2
NUM_FOODS = 50
NUM_USERS = 1000
NUM_RECORDS = 100000

foods = (list(range(NUM_FOODS)) * (math.ceil(NUM_RECORDS/NUM_FOODS)))[:NUM_RECORDS]
users = (list(range(NUM_USERS)) * (math.ceil(NUM_RECORDS/NUM_USERS)))[:NUM_RECORDS]

shuffle(foods)
shuffle(users)

df = pd.DataFrame({"user": users, "food": foods})

food_pairs = pd.Series([*itertools.combinations(df.food.unique(), SIZE_OF_PAIRS)])

start = time.time()

food_to_users = {food: set(df.user[df.food == food].unique()) for food in df.food.unique()}
out = pd.DataFrame(
    ((*pair, len(set.intersection(*(food_to_users[food] for food in pair)))) for pair in food_pairs),
    columns=[*["food" + str(i) for i in range(SIZE_OF_PAIRS)], "num users"]
)

print("Time taken: {}s".format(time.time() - start))

Upvotes: 1

jezrael
jezrael

Reputation: 862671

You can use get_dummies first:

df = pd.get_dummies(df.set_index('user'), prefix='', prefix_sep='').max(level=0)
print (df)
      candy  carrot  meat
user                     
A         1       1     1
B         0       1     1
C         0       1     1

And then count by list comprehension:

from  itertools import combinations

L = [(x[0], x[1],(df[list(x)] == 1).all(1).sum()) for x in list(combinations(df.columns, 2))]
print (L)
[('candy', 'carrot', 1), ('candy', 'meat', 1), ('carrot', 'meat', 3)]

df = pd.DataFrame(L, columns=['food 1','food 2','num users'])
print (df)
   food 1  food 2  num users
0   candy  carrot          1
1   candy    meat          1
2  carrot    meat          3

Upvotes: 3

Related Questions