How to find the total length of a column value that has multiple values in different rows for another column BY LOCATION

Viewed 32

This is part 2 question of the original problem.

Is there a way to find IDs that have both Apple and Strawberry, and then find the total length? and IDs that has only Apple, and IDS that has only Strawberry? BASED ON LOCATION

df:

        ID           Fruit        Location
0       ABC          Apple        NY            <-ABC has Apple and Strawberry
1       ABC          Strawberry   NY            <-ABC has Apple and Strawberry
2       EFG          Apple        LA            <-EFG has Apple only
3       XYZ          Apple        HOUSTON       <-XYZ has Apple and Strawberry
4       XYZ          Strawberry   HOUSTON       <-XYZ has Apple and Strawberry 
5       CDF          Strawberry   BOSTON        <-CDF has Strawberry
6       AAA          Apple        CHICAGO       <-AAA has Apple only

Desired output:

IDs that has Apple and Strawberry:
NY       1
HOUSTON  1
IDs that has Apple only:
LA       1
CHICAGO  1
IDs that has Strawberry only:
BOSTON   1

The previous code was:

v = ['Apple','Strawberry']
out = df.groupby('ID')['Fruit'].apply(lambda x: set(x) == set(v)).sum()
print(out)
>>> 2

I tried the following but it did not work and gave me the same results

v = ['Apple','Strawberry']
out = df.groupby('ID', 'LOCATION')['Fruit'].apply(lambda x: set(x) == set(v)).sum()
print(out)
>>> 2

Thanks!

1 Answers

Inefficient solution using groupby and apply

x = df.groupby('ID').agg({ 'Fruit': lambda x: tuple(x), 'Location': 'first'})
y=x.groupby('Fruit')['Location'].value_counts()

y:

Fruit                Location
(Apple,)             CHICAGO     1
                     LA          1
(Apple, Strawberry)  HOUSTON     1
                     NY          1
(Strawberry,)        BOSTON      1
Name: Location, dtype: int64

for index in set(y.index.get_level_values(0)):
    if len(index)==2:
        print(f"IDs that has {index[0]} and {index[1]}:")
        print(y.loc[index].to_string())
    else:
        print(f"IDs that has {index[0]} only:")
        print(y.loc[index].to_string())

IDs that has Apple only:
Location
CHICAGO    1
LA         1
IDs that has Apple and Strawberry:
Location
HOUSTON    1
NY         1
IDs that has Strawberry only:
Location
BOSTON    1
Related