python ziping lists of single element or tuples into list of tuples-CodePudding

Assuming I have a colection of lists of the same size (same number of elements), where the elements are single elements (strings, numbers) or tuples:

a = ['AA', 'BB', 'CC']
b = [7,8,9]
d = [('TT', 'ZZ'),('UU', 'VV'),('JJ','KK')]
e = [('mm', 'nn'), ('bb', 'vv'), ('uu', 'oo')]

I would like a way of combining whatever two lists (ideally whatever number of lists) in a way that the result is a list of tuples

# none of this works:
print(list(zip(a,b))) # this ONLY works for the case in which a & b have single elements, not tuples
print(list(zip(b,*d)))
print(list(zip(*d,*e)))

Desired result:

pseudo code:

combination of a,b = [('AA', 7), ('BB', 8), ('CC', 9)]
combination of a,d = [('AA','TT', 'ZZ'),('BB', 'UU', 'VV'),('CC','JJ','KK')]
combination of d,e = [('TT', 'ZZ','mm', 'nn'),('UU', 'VV','bb', 'vv'),('JJ','KK','uu', 'oo')]

basically the method would get the elements of the input lists treat them as tuples (even if being of only one element) and add up all the values of the tuples of the corresponding same position.

CodePudding user response：

Just flatten the output of zip, as below:

from collections.abc import Iterable

a = ['AA', 'BB', 'CC']
b = [7, 8, 9]
d = [('TT', 'ZZ'), ('UU', 'VV'), ('JJ', 'KK')]
e = [('mm', 'nn'), ('bb', 'vv'), ('uu', 'oo')]


def flatten(*e):
    for ei in e:
        yield from (ei,) if not (isinstance(ei, Iterable) or isinstance(ei, str)) else ei


print([list(flatten(ei)) for ei in zip(a, b)])
print([list(flatten(ei)) for ei in zip(a, d)])
print([list(flatten(ei)) for ei in zip(d, e)])

Output

[['AA', 7], ['BB', 8], ['CC', 9]]
[['AA', 'TT', 'ZZ'], ['BB', 'UU', 'VV'], ['CC', 'JJ', 'KK']]
[['TT', 'ZZ', 'mm', 'nn'], ['UU', 'VV', 'bb', 'vv'], ['JJ', 'KK', 'uu', 'oo']]

CodePudding user response：

Conversion to a tuple wrt. string is cumbersome in python, you have to do it like this:

def to_tuple(s):
    if isinstance(s, str):
        return s,
    try:
        return tuple(s)
    except TypeError:
        return s,

Once this is done, the rest is easy:

a_plus_d = [to_tuple(x)   to_tuple(y)
            for x, y in zip(a, d)]

CodePudding user response：

def foo(x, y):
    x = [x] if not isinstance(x, (tuple, list)) else x
    y = [y] if not isinstance(y, (tuple, list)) else y
    return [el for item in [x, y] for el in item]
    
[tuple(foo(el1, el2)) for el1, el2 in list(zip(a, d))]
# [('AA', 'TT', 'ZZ'), ('BB', 'UU', 'VV'), ('CC', 'JJ', 'KK')]

CodePudding user response：

Another short approach using itertools.chain method:

import itertools

def combine_lists(*lists):
    return list(tuple(itertools.chain(*[i if isinstance(i, (list, tuple)) else (i,) for i in items]))
                for items in zip(*lists))

print(combine_lists(a, b, e))  # [('AA', 7, 'mm', 'nn'), ('BB', 8, 'bb', 'vv'), ('CC', 9, 'uu', 'oo')]