将函数映射到子列表python的一部分

Question

I am working in python and I have a list of lists (or could be also an numpy array) like this:

list = [[0, 0, -1], [1, 1, -2], [1, 2, -3], [2, 2, -4]]

I have a function that takes for every sublist the first two elements and gives a result. Let's assume that the function I am taking about is sum() so what I want to do is

list1 = [[sum(0, 0), -1], [sum(1, 1), -2], [sum(1, 2), -3], [sum(2, 2), -4]]

and output is

list1 = [[0, -1], [2, -2], [3, -3], [4, -4]]

Is there a way to this with map or numpy.vectorize ? I can already do it as a loop but it is too expensive for my real data.

Any help, advice?

Answer 1

In case of numpy array you should have the fastest results using vectorization:

In [128]: arr = np.asarray(list)
#Out[128]:
#array([[ 0,  0, -1],
#       [ 1,  1, -2],
#       [ 1,  2, -3],
#       [ 2,  2, -4]])

In [129]: np.vstack((arr[:,0]+arr[:,1], arr[:,2])).T
#Out[129]:
#array([[ 0, -1],
#       [ 2, -2],
#       [ 3, -3],
#       [ 4, -4]])

Answer 2

Could that be

[(sum(x, y), z) for x, y, z in list]

Or if you insist on using map use itertools.starmap :

starmap( lambda x, y, z: (x+y, z), list)

?

Answer 3

For your specific example, you could use a list comprehension as follows:

my_list = [[0, 0, -1], [1, 1, -2], [1, 2, -3], [2, 2, -4]]

answer = [[x + y, z] for x, y, z in my_list]
print(answer)

Output

[[0, -1], [2, -2], [3, -3], [4, -4]]

To apply a more complex generic function to each item in your list, use map :

def func(sublist):
    return [sum(sublist[:2]), sublist[-1]]

my_list = [[0, 0, -1], [1, 1, -2], [1, 2, -3], [2, 2, -4]]

answer = map(func, my_list)

By using the function func you can customise your function as desired.