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. У кого-нибудь есть идея?

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]]