mr.bjerre
mr.bjerre

Reputation: 2898

Conditional grouping in pandas data frame

Imagine a pandas data frame given by

df = pd.DataFrame({
    'id': range(5),
    'desc': ('This is text', 'John Doe ABC', 'John Doe', 'Something JKL', 'Something more'),
    'mfr': ('ABC', 'DEF', 'DEF', 'GHI', 'JKL')
})

which yields

   id            desc  mfr
0   0    This is text  ABC
1   1    John Doe ABC  DEF
2   2        John Doe  DEF
3   3   Something JKL  GHI
4   4  Something more  JKL

I wish to determine which id's belong to eachother. Either they are matched by mfrcolumn or if mfrvalue are contained in desccolumn. E.g. id = 1 and 2 are the same group because mfr are equal but id = 0 and 1 are also the same group since ABC from mfr in id = 0 are part of desc in id = 1.

The resulting data frame should be

   id            desc  mfr  group
0   0    This is text  ABC      0
1   1    John Doe ABC  DEF      0
2   2        John Doe  DEF      0
3   3   Something JKL  GHI      1
4   4  Something more  JKL      1

Are there anyone out there with a good solution for this? I imagine that there are no really simple ones so any is welcome.

Upvotes: 2

Views: 341

Answers (1)

gyoza
gyoza

Reputation: 2152

I'm assuming 'desc' does not contain multiple 'mfr' values

Solution1:

import numpy as np
import pandas as pd

# original dataframe
df = pd.DataFrame({
    'id': range(5),
    'desc': ('This is text', 'John Doe ABC', 'John Doe', 'Something JKL', 'Something more'),
    'mfr': ('ABC', 'DEF', 'DEF', 'GHI', 'JKL')
})

# for final merge
ori = df.copy()

# max words used in 'desc'
max_len = max(df.desc.apply(lambda x: len(x.split(' '))))

# unique 'mfr' values
uniq_mfr = df.mfr.unique().tolist()

# if list is less than max len, then pad with nan
def padding(lst, mx):
    for i in range(mx):
        if len(lst) < mx:
            lst.append(np.nan)
    return lst
df['desc'] = df.desc.apply(lambda x: x.split(' ')).apply(padding, args=(max_len,))

# each word makes 1 column
for i in range(max_len):
    newcol = 'desc{}'.format(i)
    df[newcol] = df.desc.apply(lambda x: x[i])
    df.loc[~df[newcol].isin(uniq_mfr), newcol] = np.nan

# merge created columns into 1 by taking 'mfr' values only
df['desc'] = df[df.columns[3:]].fillna('').sum(axis=1).replace('', np.nan)

# create [ABC, ABC] type of column by merging two columns (desc & mfr)
df = df[df.columns[:3]]
df.desc.fillna(df.mfr, inplace=True)
df.desc = [[x, y] for x, y in zip(df.desc.tolist(), df.mfr.tolist())]
df = df[['id', 'desc']]
df = df.sort_values('desc').reset_index(drop=True)

# BELOW IS COMMON WITH SOLUTION2
# from here I borrowed the solution by @mimomu from below URL (slightly modified)
# try to get merged tuple based on the common elements
# https://stackoverflow.com/questions/4842613/merge-lists-that-share-common-elements
import itertools

L = df.desc.tolist()
LL = set(itertools.chain.from_iterable(L)) 

for each in LL:
    components = [x for x in L if each in x]
    for i in components:
        L.remove(i)
        L += [tuple(set(itertools.chain.from_iterable(components)))]

# allocate merged tuple to 'desc'
df['desc'] = sorted(L)

# grouping by 'desc' value (tuple can be key list cannot be fyi...)
df['group'] = df.groupby('desc').grouper.group_info[0]

# merge with the original
df = df.drop('desc', axis=1).merge(ori, on='id', how='left')
df = df[['id', 'desc', 'mfr', 'group']]

Solution2 (2nd half is common with Solution1):

import numpy as np
import pandas as pd

# original dataframe
df = pd.DataFrame({
    'id': range(5),
    'desc': ('This is text', 'John Doe ABC', 'John Doe', 'Something JKL', 'Something more'),
    'mfr': ('ABC', 'DEF', 'DEF', 'GHI', 'JKL')
})

# for final merge
ori = df.copy()

# unique 'mfr' values
uniq_mfr = df.mfr.unique().tolist()

# make desc entries as lists
df['desc'] = df.desc.apply(lambda x: x.split(' '))

# pick up mfr values in desc column otherwise nan
mfr_in_descs = []
for ds, ms in zip(df.desc, df.mfr):
    for i, d in enumerate(ds):
        if d in uniq_mfr:
            mfr_in_descs.append(d)
            continue
        if i == (len(ds) - 1):
            mfr_in_descs.append(np.nan)

# create column whose element is like [ABC, ABC]
df['desc'] = mfr_in_descs
df['desc'].fillna(df.mfr, inplace=True)
df['desc'] = [[x, y] for x, y in zip(df.desc.tolist(), df.mfr.tolist())]
df = df[['id', 'desc']]
df = df.sort_values('desc').reset_index(drop=True)

# BELOW IS COMMON WITH SOLUTION1
# from here I borrowed the solution by @mimomu from below URL (slightly modified)
# try to get merged tuple based on the common elements
# https://stackoverflow.com/questions/4842613/merge-lists-that-share-common-elements
import itertools

L = df.desc.tolist()
LL = set(itertools.chain.from_iterable(L)) 

for each in LL:
    components = [x for x in L if each in x]
    for i in components:
        L.remove(i)
        L += [tuple(set(itertools.chain.from_iterable(components)))]

# allocate merged tuple to 'desc'
df['desc'] = sorted(L)

# grouping by 'desc' value (tuple can be key list cannot be fyi...)
df['group'] = df.groupby('desc').grouper.group_info[0]

# merge with the original
df = df.drop('desc', axis=1).merge(ori, on='id', how='left')
df = df[['id', 'desc', 'mfr', 'group']]

From 2 solutions above, I get the same results df:

    id  desc            mfr  group
0   0   This is text    ABC  0
1   1   John Doe ABC    DEF  0
2   2   John Doe        DEF  0
3   3   Something JKL   GHI  1
4   4   Something more  JKL  1

Upvotes: 1

Related Questions