How can I vectorize this loop? (Not how do I one-hot encode)

Viewed 56

How can this loop be vectorized? Or can it? I just don't see it.

for c in cats:
    df[c] = np.where(df['categories'].str.contains(c),1,0)   

(Before reading the full story below, expect that you will be very tempted to answer what appears to be the a one-hot encoding question with get_dummies(sep=';). But that is not really what I am asking. What I strictly want answered is the question above - about vectorizing that loop. I included the full story for context only.)

The full story:

d = {'categories': ['catA;catB;catC','catA;catB','catC','catB;catC']}
df = pd.DataFrame(d)

print(df)

       categories
0  catA;catB;catC
1       catA;catB
2            catC
3       catB;catC

cats = ['catA','catB','catC']

for c in cats:
    df[c] = np.where(df['categories'].str.contains(c),1,0)
    
print(df)

       categories  catA  catB  catC
0  catA;catB;catC     1     1     1
1       catA;catB     1     1     0
2            catC     0     0     1
3       catB;catC     0     1     1
1 Answers

Based on what I understood from the question and gathered from comments. you are not looking for get_dummies since you may have a larger string not necessarily seperated by the ;. In that case its better that you use series.str.findall to return the values which match strings in cat. Then you can either join them and use str.get_dummies() or MultiLabelBinarizer:

cats = ['catA','catB','catC']
s=df['categories'].str.findall('|'.join(cats))

Fast method:

from sklearn.preprocessing import MultiLabelBinarizer
mlb = MultiLabelBinarizer()
out1 = df.join(pd.DataFrame(mlb.fit_transform(s),
         columns=mlb.classes_, index=df.index))

using series.str.get_dummies()

out2 = df.join(s.str.join("|").str.get_dummies())

print((out1==out2).all().all())
#True

print(out1)

       categories  catA  catB  catC
0  catA;catB;catC     1     1     1
1       catA;catB     1     1     0
2            catC     0     0     1
3       catB;catC     0     1     1
Related