numpy 向量化在 (2,) 数组上的使用

Question

I have a numpy array of (m, 2) and I want to transform it to shape of (m, 1) using a function below.

def func(x):
    if x == [1., 1.]:
        return 0.
    if x == [-1., 1.] or x == [-1., -1.]:
        return 1.
    if x == [1., -1.]:
        return 2.

I want this for applied on each (2,) vector inside the (m, 2) array resulting an (m, 1) array. I tried to use numpy.vectorize but it seems that the function gets applied in each element of a array (which makes sense in general purpose case). So I have failed to apply it.

My intension is not to use for loop. Can anyone help me with this? Thanks.

Answer 1

import numpy as np


def f(a, b):
    return a + b


F = np.vectorize(f)
x = np.asarray([[1, 2], [3, 4], [5, 6]]).T

print(F(*x))

Output:

[3, 7, 11]