如何优雅地用字典dict映射numpy array：map和np.vectorize

问题

现在有一个ndarray和字典dict如下：

a1 = np.array(['a', 'b'])
d = {'a':1, 'b':2}

想根据字典d中的对应关系将a1映射成如下：

array([1, 2])

方案一：map函数映射一维ndarray

第一反应是使用map函数，不过其中需要注意的是，map中传入的函数部分需为dict.get

a1 = np.array(['a', 'b'])
d = {'a':1, 'b':2}
a2 = np.array(list(map(d.get, a1)))
a2
>>> array([1, 2])

如上，一行代码就解决了，缺点是只可以用在一维ndarray上。

方案二：np.vectorize映射多维ndarray

对于多维ndarray，可以使用numpy的vectorzie函数。

a1 = np.array([['a', 'b'], ['b', 'a']])
d = {'a':1, 'b':2}
a2 = np.vectorize(d.get)(a1)
a2
>>> array([[1, 2],
       [2, 1]])

依旧一行代码解决问题，优雅！