python - Unravel Index numpy - 自己的实现

Question

我尝试自己np.unravel_index实施np.ravel_multi_index。因为np.ravel_multi_index我可以编写这个简短的函数：

def coord2index(coord, shape):
    return np.concatenate((np.asarray(shape[1:])[::-1].cumprod()[::-1],[1])).dot(coord) 

但我很难为np.unravel_index. 有人有想法吗？

标签： pythonnumpyindexing

Answer 1

这是一种可能的实现：

import numpy as np

def index2coord(index, shape):
    return ((np.expand_dims(index, 1) // np.r_[1, shape[:0:-1]].cumprod()[::-1]) % shape).T

shape = (2, 3, 4)
coord = [[0, 1], [2, 0], [1, 3]]
print(index2coord(coord2index(coord, shape), shape))
# [[0 1]
#  [2 0]
#  [1 3]]