How to pythonically iterate two lists in an arbitrary order

Viewed 190

I want to iterate over two lists in such a way that I can take an arbitrary number of values from one list and maintain my position in the other.

I've used indexes to store the current position in each list, then a while loop to go over them, but this definitely isn't very pythonic.

def alternate_iterate(a,b,cond=lambda x, y : x > y):
    pos_a = 0
    pos_b = 0
    retval = []

    while(True):

        if(pos_a == len(a) and pos_b == len(b)):
            break

        if(pos_a < len(a) and cond(a[pos_a],b[pos_b])):
            retval += [a[pos_a]]
            pos_a += 1
        elif(pos_b < len(b)):
            retval += [b[pos_b]]
            pos_b += 1

    return retval

#example usage
print(alternate_iterate(['abc','abcd','ab','abc','ab'],
                        ['xy','xyz','x','xyz'],
                        cond=lambda x,y: len(x) > len(y))

This should print ['abc','abdc','xy','xyz','ab','abc','ab','x','xyz'], where you don't have a perfect 1:1 alternating order. The order of the elements and the type of the elements should only depend on whatever cond is defined as.

4 Answers

The more Pythonic way is usually to not use indexes at all and it is preferable not to use exceptions as a means of controlling "intended" program logic. You should also avoid unnecessary parentheses.

Here's how you could do it using iterators:

def merge(a, b, cond=lambda x, y : x < y):
    Done           = []
    iterA, iterB   = iter(a), iter(b)
    valueA, valueB = next(iterA, Done), next(iterB, Done)
    result         = []
    while not(valueB is Done and valueA is Done):
        if valueB is Done or valueA is not Done and cond(valueA, valueB):
            result.append(valueA)
            valueA = next(iterA, Done)
        else:
            result.append(valueB)
            valueB = next(iterB, Done)
    return result

This has the added benefit of making the function work efficiently with any iterable data as parameters.

for example:

print(merge(range(5, 10), range(7, 15)))

# [5, 6, 7, 7, 8, 8, 9, 9, 10, 11, 12, 13, 14]

It also makes it easy to create an iterator version of the function for lazy traversal:

def iMerge(a, b, cond=lambda x, y : x < y):
    Done           = []
    iterA, iterB   = iter(a), iter(b)
    valueA, valueB = next(iterA, Done), next(iterB, Done)
    while not(valueB is Done and valueA is Done):
        if valueB is Done or valueA is not Done and cond(valueA, valueB):
            yield valueA
            valueA = next(iterA ,Done)
        else:
            yield valueB
            valueB = next(iterB, Done)

EDIT Changed None to Done in order to let the function support None as a legitimate value in the input lists.

Welcome to Stackoverflow. Summarising, you appear to want to take a value from one list or the other depending on the value of some predicate. Your existing logic doesn't appear to take into account the possibility of one of the lists being exhausted, at which point I've assume you would want to copy any remaining values from the other list.

Rather than using index values to select the successive list elements you can build an iterator on the list, and use the next function to get the next value.

In that case your logic would end up looking something like this:

def alternate_iterate(a_lst, b_lst, cond=lambda x, y: x > y):
    a_iter = iter(a_lst)
    b_iter = iter(b_lst)
    a = next(a_iter)
    b = next(b_iter)
    ret = []
    while True:
        if cond(a, b):
            ret.append(a)
            try:
                a = next(a_iter)
            except StopIteration:
                ret.append(b)
                for x in b_iter:
                    ret.append(x)
                return ret
        else:
            ret.append(b)
            try:
                b = next(b_iter)
            except StopIteration:
                ret.append(a)
                for x in a_iter:
                    ret.append(x)
                return ret


print(alternate_iterate(['abc','abcd','ab','abc','ab'],
                        ['xy','xyz','x','xyz'],
                        cond=lambda x,y: len(x) > len(y)))

the result I get is

['abc', 'abcd', 'xy', 'xyz', 'ab', 'abc', 'ab', 'x', 'xyz']

which would appear to be what you would expect.

As is often the case in such examples, you write more logic to handle the rarer corner cases (in this case, one list or the other becoming exhausted) than you do to handle the "happy path" where things proceed as normal.

This version is using only iterators to achieve the functionality lazilly (which is Pythonic):

a = ['abc','abcd','ab','abc','ab']
b = ['xy','xyz','x','xyz']

cond=lambda x,y: len(x) > len(y)

def alternate_iterate(a, b, cond):
    a, b = iter(a), iter(b)

    def _return_rest():
        def _f(val, it):
            yield val
            yield from it
        return _f

    v1, v2 = next(a, _return_rest), next(b, _return_rest)

    while True:
        if v1 is _return_rest:
            yield from v1()(v2, b)
            break

        if v2 is _return_rest:
            yield from v2()(v1, a)
            break

        if cond(v1, v2):
            yield v1
            v1 =  next(a, _return_rest)
        else:
            yield v2
            v2 = next(b, _return_rest)

print(list(alternate_iterate(a, b, cond)))

Prints:

['abc', 'abcd', 'xy', 'xyz', 'ab', 'abc', 'ab', 'x', 'xyz']

Put your lists into generators, and then you can call next on each one to get the next value. This answer isn't meant to be a complete solution, just to show how generators can produce values in any order with very simple, Pythonic code:

agen = iter(a)
bgen = iter(b)
print next(agen) # 'a'
print next(bgen) # 1
print next(bgen) # 2
print next(agen) # 'b'

and so on.

Related