donopj2
donopj2

Reputation: 4123

Generator to Yield a List

I have three lists of lists, and I'm trying to write a generator function to help me package up values in the same index.

So my lists:

list1 = [[1, 2, 3], [2, 3, 4],...]
list2 = [[4, 5, 6], [5, 6, 7],...]
list3 = [[8, 9, 10], [9, 10, 11],...]

My desired output:

result1 = [[1, 4, 8], [2, 5, 9],...]
result2 = [[2, 5, 9], [3, 6, 10],...]
result3 = [[3, 6, 10], [4, 7, 11],...]

My attempt:

def bundle(x, y, z, index):
    for row in x, y, z:
        for item in row[index]:
            yield list(item)

I keep getting float is not iterable errors. If i modify it slightly:

def bundle(x, y, z, index):
    for row in x, y, z:
        for item in row:
            yield item[index]

I get the values I want as one large sequence, but I would prefer to keep them grouped in a nested style.

Upvotes: 1

Views: 2121

Answers (3)

Will Furnass
Will Furnass

Reputation: 108

If you're dealing with numeric values then you could use numpy's transposition method to achieve what you want:

import numpy
numpy.array([list1,list2, list3]).T

Upvotes: 2

Niklas B.
Niklas B.

Reputation: 95358

If you're dealing with large lists, a custom, fully lazy approach would be the following:

import itertools as it

def bundle(lists, index):
  return ([b[index] for b in blocks] for blocks in it.izip(*lists))

print list(bundle([[[1, 2, 3],  [2, 3, 4]], 
                   [[4, 5, 6],  [5, 6, 7]], 
                   [[8, 9, 10], [9, 10, 11]]], 
                  0))
# => [[1, 4, 8], [2, 5, 9]]

Upvotes: 2

Sven Marnach
Sven Marnach

Reputation: 602705

One way to do this is by a repeated application of zip():

res1, res2, res3 = zip(*(zip(*x) for x in zip(list1, list2, list3)))

This uses zip(list1, list2, list3) to create a sequence of matrices, zip(*x) to transpose each of this matrices, and a final zip() to unpack to the three resulting sequences. (I don't think this approach is very efficient.)

Upvotes: 2

Related Questions