CODEWITHSUNDEEP
pythonarraysnumpy
Stat_prob_001
Stat_prob_001

Reputation: 181

numpy vectorize use on (2,) array

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.

Upvotes: 0

Views: 156

Answers (1)

Lorenzo Palloni
Lorenzo Palloni

Reputation: 54

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]

Upvotes: 1

Related Questions