How to combine elements of two uneven lists in python

I have three lists:

proteins = ["1ab1_P.xyz", "1wav_P.xyz"]
ligands = ["DMSO_L.xyz", "MOH_L.xyz", "EOH_L.xyz"]
results = [5.2, 5.0, 6.7, 4.1, 3.9, 0.22]

I want to get a result like:

output = [
    '1ab1DMSO5.2',
    '1ab1MOH5.0',
    '1ab1EOH6.7',
    '1wavDMSO4.1',
    '1wavMOH3.9',
    '1wavEOH0.22'
]

The idea is that each element could be computed like

def an_output(protein, ligand, result):
    return protein.split('_')[0] + ligand.split('_')[0] + str(result)

The pattern is that each protein should be matched up with each ligand in a specific order (3 * 2 = 6 pairs), and each of these matched up with a corresponding result in sequence.

I have found that the code

from itertools import cycle     
aa = [ele for comb in (zip(cycle(list1), list2) for ele in comb]

can help with simple lists, but I can't implement it because I need to use split as well in the elements of my lists.

How I can achieve the desired result regardless of the lists' length?

Upvotes: 0

Views: 62

Answers (1)

Adon Bilivit
Adon Bilivit

Reputation: 26925

I think this may be what you're looking for...

from itertools import product
from functools import cache

proteins = ["1ab1_P.xyz", "1wav_P.xyz"]
ligands = ["DMSO_L.xyz", "MOH_L.xyz", "EOH_L.xyz"]
results = [5.2, 5.0, 6.7, 4.1, 3.9, 0.22]

@cache
def prefix(s):
    return s.split('_')[0]

output = []

for (p, l), r in zip(product(proteins, ligands), results):
    output.append(f'{prefix(p)} {prefix(l)} {r}')

print(output)

Output:

['1ab1 DMSO 5.2', '1ab1 MOH 5.0', '1ab1 EOH 6.7', '1wav DMSO 4.1', '1wav MOH 3.9', '1wav EOH 0.22']

Note:

Spaces inserted into output for clarity

Upvotes: 1

Related Questions