简体   繁体   中英

How to create a numpy 2-D array whose elements depend on elements of two lists already fixed

Example:

Let

A = np.array([1,2,3,5,7])
B = np.array([11,13,17,19,23]) 

and I would like to create a matrix C of size (5,5) whose elements are

c_ij = f(a[i],b[j])

where f is a fixed function for example f(x,y) = x*y + x + y which means

c_ij = a[i]*b[j] + a[i] + b[j]

In the case where c_ij depends on i and j only and does not depend on the lists A and B , we can use np.fromfunction(lambda i,j: f(i,j), (5,5)) but it's not the case.

I would like to know how we can do that?

import numpy as np
# set up f()
def _my_math_function(x, y):
    return x*y + x + y

# variable setup
A = np.array([1,2,3,5,7])
B = np.array([11,13,17,19,23]) 

# nested comprehensive loop
# basically f(1,11), (1,13) ... f(7,19) f(7,23)
c = [_my_math_function(a,b) for b in B for a in A]    

# len list
shape_a = len(A)
shape_b = len(B)
c = np.array(c).reshape(shape_a,shape_b)

# Results of c
array([[ 23,  35,  47,  71,  95],
       [ 27,  41,  55,  83, 111],
       [ 35,  53,  71, 107, 143],
       [ 39,  59,  79, 119, 159],
       [ 47,  71,  95, 143, 191]])

Is this what you want:

def bar(arr1, arr2, func):
    ind1, ind2 = np.meshgrid(range(len(arr1)), range(len(arr2)))
    x = arr1[ind1]
    y = arr2[ind2]
    return func(x, y)

bar(A, B, f)

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM