Count rows and create label column considering the position of the string [with str.contains] in python

Viewed 46

I have the following dataframe:

df = pd.DataFrame([["1", "shirt is blue", True, "Feb"],
                   ["2", "blue shirt with green stripes", False, "March"],
                   ["3", "shirt is green", True, "Feb"],
                   ["4", "shirt is light-green", True, "Feb"],
                   ["5", "shirt is light-green", True, "Feb"],
                   ["6", "blue shirt", True, "March"],
                   ["7", "shirt shirt shirt", False, "March"],
                   ["8", "green shirt with dark-blue stripes", False, "Feb"],
                   ["9", "shirt is blue with red stripes", False, "March"],
                   ["10", "shirt is a red shirt", False, "Feb"]],
                  columns=["id", "text", "cond", "month"])

I'd like to count (to plot a graph later) how many are the shirts of each color (blue, green and red).

Tried a filter like this:

 blue = df[df['text'].str.contains('blue',case=False)]
 green = df[df['text'].str.contains('green',case=False)]
 red = df[df['text'].str.contains('red',case=False)]

First question: Is there a smarter way to do this (as in a loop, instead of a line for each color)?

Second (and more important) question: I should consider the "blue shirt with green stripes" as a blue shirt, but when I use the code above, the order of the "colors" in the sentence does not count at all, so I get one count for blue and other for green with this text.

How can I get one count only for blue, in cases like the example?

It's important to notice that the first color that appears in the sentence is always the "predominant" one.

So, how can I do it properly, that is, count the rows considering the position of the string? (Or maybe shoud I consider the intersections?)

After solving the above, I also want to add a new column:

df['color'] = 'NaN'
df['color'] = np.where((df['id'].isin(blue['id'])),'Blue',df.color)
df['color'] = np.where((df['id'].isin(green['id'])),'Green',df.color)
df['color'] = np.where((df['id'].isin(red['id'])),'Red',df.color)
df['color'] = df['color'].str.replace('NaN', 'Other', regex=True)

As in the first question, any way to do this in a smarter way?

1 Answers

If first color is "predominant", you can use .str.extract. It will extract first match from the text:

df["color"] = df["text"].str.extract(r"(red|green|blue)")
print(df)

Prints:

   id                                text   cond  month  color
0   1                       shirt is blue   True    Feb   blue
1   2       blue shirt with green stripes  False  March   blue
2   3                      shirt is green   True    Feb  green
3   4                shirt is light-green   True    Feb  green
4   5                shirt is light-green   True    Feb  green
5   6                          blue shirt   True  March   blue
6   7                   shirt shirt shirt  False  March    NaN
7   8  green shirt with dark-blue stripes  False    Feb  green
8   9      shirt is blue with red stripes  False  March   blue
9  10                shirt is a red shirt  False    Feb    red
Related