如何将依赖输入位置的函数映射至NumPy数组？

这个场景我之前也碰到过！当函数需要用到元素在数组里的位置索引时，直接对数组元素映射就行不通了，你目前用meshgrid的思路是对的，但完全可以去掉循环，用NumPy的向量化特性来实现更高效、更简洁的代码。

下面分享几种规范且高效的实现方案：

1. 优先选择：np.indices + 广播机制（性能最优） #

NumPy的np.indices()可以直接生成和原数组同形状的坐标矩阵，不用手动做meshgrid再调整维度，然后利用广播特性直接完成元素级运算，全程无循环，完全依托NumPy的底层优化，性能拉满。

示例代码：

import numpy as np

# 生成测试数组
arr = np.random.rand(3, 3)

# 生成对应每个元素的行索引(x)和列索引(y)矩阵
x, y = np.indices(arr.shape)

# 直接应用依赖位置的函数（这里以x*y*val为例）
result = x * y * arr

这里x是形状为(3,3)的矩阵，每个位置的值是该元素的行索引；y同理是列索引矩阵。直接做算术运算时，NumPy会自动广播维度，完成所有元素的计算，比循环快几个数量级。

2. 复杂逻辑场景：np.vectorize（代码简洁） #

如果你的函数逻辑比较复杂（不是简单的算术组合），可以用np.vectorize()将普通函数转化为能处理数组的向量化函数，代码会更简洁。不过要注意：np.vectorize本质是伪向量化，内部还是循环，性能不如纯向量化操作，但胜在代码易读。

示例代码：

import numpy as np

def custom_func(x, y, val):
    # 这里可以写任意复杂的逻辑，比如结合索引和值的条件判断
    if x + y > 2:
        return x * y * val ** 2
    else:
        return val + np.cos(x)

arr = np.random.rand(3, 3)
x, y = np.indices(arr.shape)

# 将自定义函数向量化，直接传入坐标矩阵和原数组
result = np.vectorize(custom_func)(x, y, arr)

3. 避免使用：循环遍历（性能最差） #

你之前用ravel加循环的方式，虽然能实现功能，但在数组规模较大时，Python循环的开销会非常明显，完全不建议在生产代码中使用。NumPy的核心优势就是向量化操作，能不用循环就尽量不用。

总结 #

• 如果是简单的算术运算，优先用np.indices+广播，这是最规范、性能最高的方案；
• 如果是复杂逻辑，用np.vectorize简化代码，但如果追求极致性能，建议把逻辑拆解成向量化操作；
• 绝对避免手动循环遍历数组元素，尤其是大数组场景。

内容的提问来源于stack exchange，提问作者C4rnie