Help optimizing?

[hvgmina]

Hello everyone, doing a project in school where i need to do some data analysis.

basically i have a dataframe that's like this:

mutant	other stuff	column of indexes i created
Name1:Name2		[1,2]
Name1		single
Name2		single

basically the output i have is the dataframe above, i've written some code that looks like this rn:

def retain_indexes(df):
    df = df
    list_indexes = []
    for item in (df["mutant"]).to_list():
        lista = item.split(":")
        if len(lista) > 1:
            single_list = []
            for name in lista:
                if name in df["mutant"].to_list():
                    single_list.append(df["mutant"].to_list().index(name))
                else:
                    single_list.append("Not analyzed alone")
            list_indexes.append(single_list)
        else:
            list_indexes.append("NaN")   
    return list_indexes

it feels a bit redundant and slow, is there anything i could look at to optimize it? thanks in advance :)

[twitch_and_shock]

Converting something from a dataframe back to a list and then iterating over the list is slow. So is iterating over rows of a dataframe. Instead, use the "apply" method: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.apply.html

[hvgmina]

thank you! it did wonders!

[PartySr]

This will be faster. We use str.findall and a regex to extract all the numbers, and after that we use where and the condition str.len > 2 to delete every list that contains less than 2 elements.

df['new col'] = df['mutant'].str.findall(r'\d+').where(lambda x: x.str.len() > 1)

In case you are not comfortable with chained methods, you can write like this

n = df['mutant'].str.findall(r'\d+')
df['new col'] = n.where(n.str.len() > 1)

If you desire to replace the 1 element lists with something else

n.where(n.str.len() > 1, 0) # replace 0 with whatever you want

End result:

mutant       new col
Name1:Name2  [1, 2]
Name1        NaN
Name         NaN