Filtering a List based on a Suffix and avoid duplicates

Question

I was wondering if there is a better way to solve this issue.

I have a list like this:

imp_list=["aa","daa","ab","as","aem",
      "aum","aw","aa","acm","at",
      "ar","aa_imp","daa_imp","ab_imp",
      "as_imp"]

I want to select all the strings that have the _imp suffix, plus all the the strings that have no _imp partner.

That's because aa_imp it's just a modified version of aa, and for me, in a large sense, it's a replica. So I created this function:

def imputed_handler(my_list): 
  imp=[x for x in my_list if "imp" in x]

  cleaned_imp=set(map(lambda x: x.replace("_imp",""),imp))

  not_imp=[x for x in my_list if "imp" not in x]
  set_list=set(my_list)
  no_replica=list(set_list-cleaned_imp)
  print(no_replica)
  return no_replica

Running the code as described above

test=imputed_handler(imp_list)

I get the following output:

['at', 'acm', 'aw', 'ar', 'daa_imp', 'aem', 'as_imp', 'ab_imp', 'aa_imp', 'aum']

Do better solutions exist? Thanks for your time, and let me know if something is not clear :)

Answer 1

First off, since you're looking for strings that end with _imp, x.endswith("_imp") is probably more reliable than "imp" in x.

Second, since we only want to remove _imp when it's a suffix, and since we know how long it is, we can use string slicing to remove the last 4 characters. That means we don't accidentally turn x_impabc_imp into xabc instead of x_impabc or something. Assuming that's not what we want to do. Maybe it is, maybe it isn't, maybe we'll never get an input like that so it won't matter either way, I have no idea.

I'm also thinking it might be nice to pass the suffix as a parameter. Maybe it might change at some point in the future. Maybe another part of the program uses a different suffix in a similar way. Doesn't hurt to be prepared.

In general, though, your approach is good. You find the _imp ones, you figure out what they are replacing, you remove the ones that have been replaced.

The neat thing is, you can do the first two steps all at once in a single pass through the input list. For example, you can do this:

def imputed_handler(my_list, suffix="_imp"):
    my_set = set(my_list)
    cleaned_imp = set(item[:-len(suffix)] for item in my_set if item.endswith(suffix))

    return list(my_set - cleaned_imp)

Though a loop might be clearer at that point:

def imputed_handler(my_list, suffix="_imp"):
    my_set = set(my_list)

    for item in my_list:
        if item.endswith(suffix):
            my_set.discard(item[:-len(suffix)])

    return list(my_set)

Answer 2

PEP 8

The Style Guide for Python Code lists several conventions Python programs should follow. Things like:

• spaces around binary operators (imp = [...] and set_list = set(my_list)).
• spaces after comma (eg, ""), imp))

Side Effects

You are returning a result from imputed_hander; shouldn't print from inside it.

Type Hints and Doc Strings

What is def imputed_handler(my_list)? What does it do? What is my_list a list of? What is returned?

from typing import List

def imputed_handler(my_list: List[str]) -> List[str]: 
    """
    Filter a list of strings, removing `"xyz"` if `"xyz_imp"` is found.
    Return the filtered list.
    """

Now we have a (poor) description, and can see the argument & return types.

Improved code

As mentioned in Sara J's answer, .endswith() is preferable to simply checking if the search string is contained anywhere inside the original.

Converting the list into a set, and then back into a list ruins the original lists order. Here is a possible improvement on Sara's solution:

from typing import List

def imputed_handler(my_list: List[str], suffix: str = "_imp") -> List[str]:
    """
    A good description here.
    """

    unwanted = {item.removesuffix(suffix) for item in my_list if item.endswith(suffix)}
    return [item for item in my_list if item not in unwanted]

A set is only constructed from the prefix of items which end in _imp. This should result in a smaller memory footprint, that set(my_list) - cleaned_imp. Since it is a set, the in operator is \$O(1)\$, filtering the list is fast.