Searching for many different names in a Panda frame

Question

Consider the following example code:

c1 = ["Mary herself","Mary is mine","John himself","John is mine","Alan himself","Alan is mine"]
s1 = pd.Series(c1)
c2 = [1,2,3,4,5,6]
s2 = pd.Series(c2)

dictionary = {"A":s1,"B":s2}

df = pd.DataFrame(dictionary)
print(df)
for i in range(6):
  name = df["A"].iloc[i]
  if "Mary" in name:
    print(name)
  if "Alan" in name:
    print(name)
  else:
    print("no")

I want to filter strings that include "Mary" or "Alan". My solution works well for only two cases, but I wonder if there is a more suitable solution for a lot more cases. For example, if I want to filter "John" and "Brown" too, I would already need four if statements.

How can I improve my code so I wouldn't have to add a new if statement for each name I want to filter?

Answer 1

Your code follows this pseudo-code pattern:

if any of these strings are in name:
    do x

The simplest way to express that in Python is to invert the condition:

if name in any of these strings:
    do x

Or as actual Python:

if name in ["Alan", "Mary"]:
    print(name)

Answer 2

elif

Your code doesn't look like it will work the way you intend it to:

  if "Mary" in name:
    print(name)
  if "Alan" in name:
    print(name)
  else:
    print("no")

If name contains "Mary", but does not contain "Alan", you will have two lines printed, due to the first print(name) and the print("no") statements. It would also print name twice in the string contained both "Mary" and "Alan", like "Mary is Alan's"`.

You probably wanted an if/elif/else statement:

  if "Mary" in name:
    print(name)
  elif "Alan" in name:
    print(name)
  else:
    print("no")

This will print only once per row.

.startswith

Do you really want in? Should "Anne" match "Mary-Anne is herself"? Or are you looking for name.startswith("...") tests?

Should "Mary" match "Maryanne is herself" or "Mary-Anne is herself"? Maybe you want to add a space to the end of the search term:

  if name.startswith("Mary "):
      print(name)
  elif name.startswith("Alan "):
      print(name)
  else
      print(no)

Alternately, you may want to split name into words, and check for equality. You'll have to clarify your question.

or

If you want to do the same thing with multiple conditions, you could link them with or:

  if name.startswith("Mary ") or name.startswith("Alan "):
      print(name)
  else
      print(no)

any

If you want to test a long list of similar conditions, joined together with or, you are really wanting to test if any of the conditions match. Which is perfect for the any(...) function, which returns True if any of the conditions is True.

Combined any() with a generator expression, to generate a series of conditions to test:

prefixes = ["Mary ", "Alan "]
for i in range(6):
    name = df["A"].iloc[i]
    if any(name.startwith(prefix) for prefix in prefixes):
        print(name)
    else:
        print("no")

Loop like a native

Why are you looping over indices, and then extracting the data by index? Why hard-code the length of the range?

prefixes = ["Mary ", "Alan "]
for name in df["A"]:
    if any(name.startwith(prefix) for prefix in prefixes):
        print(name)
    else:
        print("no")

Answer 3

If your series has many more rows than you have names to check (which should be the case), you should use the vectorized string functions in pandas.

names = ["Mary", "Alan"]
names_re = "|".join(names)

df = pd.DataFrame({"A": ["Mary herself","Mary is mine","John himself","John is mine","Alan himself","Alan is mine"],
                   "B": [1,2,3,4,5,6]})
df[df["A"].str.contains(names_re)]

#               A  B
# 0  Mary herself  1
# 1  Mary is mine  2
# 4  Alan himself  5
# 5  Alan is mine  6

This is because iterating over a series using Python is much slower than these vectorized functions which are run in C. Note that the combined search string is a regex looking for any of the names. Don't do this if you a hundred names or more.