Python : how to make 1-D array from 2-D array

Question

I have array2D = [[1,2,3],[4,5,6]]. What I want is a function which takes an index and returns the elements in 1D array.

Example: fn(0) -> returns [1,4]
         fn{1) -> returns [2,5]

I need a fast way to do this.

Answer 1

you can use lambda and list comprehension:

array2D = [[1,2,3],[4,5,6]]

fn = lambda x: [item[x] for item in array2D]

print(fn(0)) # [1, 4]
print(fn(1)) # [2, 5]
print(fn(2)) # [3, 6]

as suggested in the comments, you may apply the same concept with a function definition:

def fn(x): return [item[x] for item in array2D]

print(fn(0)) # [1, 4]
print(fn(1)) # [2, 5]
print(fn(2)) # [3, 6]

Lambda functions are pretty useful, and let you define operation in a really clear way.
In our example, our lambda accept a variable x, which represent the index we want of each item in array2D

Then you have list comprehension, similarly to lambda function, they are a really powerful tool and a must in python

In this situation you should prefear the function definiton, as suggested by PEP-8.

Answer 2

The following list comprehension will work:

def fn(i, lst):
    return [sublst[i] for sublst in lst]

>>> array2D = [[1, 2, 3], [4, 5, 6]]
>>> fn(0, array2D)
[1, 4]
>>> fn(1, array2D)
[2, 5]

Answer 3

You can use operator.itemgetter:

array2D = [[1,2,3],[4,5,6]]

from operator import itemgetter

def fn(x, k):
    return list(map(itemgetter(k), x))

fn(array2D, 0)  # [1, 4]

If you want to define new functions for retrieving a specific index, you can do so via functools.partial:

from functools import partial

def fn(x, k):
    return list(map(itemgetter(k), x))

get_zero_index = partial(fn, k=0)
get_zero_index(array2D)  # [1, 4]

Answer 4

How about a generator?

We could use zip to pack them, then create a empty list to store the generated data:

class myZip(object):
    __slots__ = ('zipData', 'interList')

    def __init__(self, *args):
        self.zipData = zip(*args)
        self.interList = []

    def __call__(self, index):
        try:
            return self.interList[index]
        except IndexError:
            try:
                if index == 0:
                    self.interList.append(next(self.zipData))
                    return self.interList[index]
                for i in range(index-(len(self.interList)-1)):
                    self.interList.append(next(self.zipData))
                return self.interList[index]
            except StopIteration:
                raise IndexError("index out of range")

    def __iter__(self):
        for i in self.interList:
            yield i

        for i in self.zipData:
            yield i

array2D = [[1,2,3],[4,5,6]]

a = myZip(*array2D)

print(a(2))
print(a(1))
print(a(0))
---
(3, 6)
(2, 5)
(1, 4)

The benefits of this is we do not need to produce all data at once.

Answer 5

Here are my two cents using slicing (I have to use additional np.array() for this because your original data was a list):

array2D = np.array([[1,2,3],[4,5,6]])

def fn(n): return (list(array2D[:,n]))

print (fn(0), fn(1), fn(2))