Cross-posted from here.
Question
Let $A$ be a $M$-dimensional tensor of size $m_1 \times m_2 \times \cdots \times m_M$ and $B$ be a $N$-dimensional tensor of size $n_1 \times n_2 \times \cdots \times n_N.$ Both have the same number of elements.
I have the index of one element in $A: (i_1, i_2, i_3, \cdots)$ and I want to find the index of the element with the same position in $B: (j_1, j_2,\cdots).$
What formula can I use?
The tensors are zero-indexed and row-major. The sequence of elements is preserved.
Answer
The easiest way to do this is to imagine mapping from $A$ to a $1$-dimensional tensor, call it $C,$ and then then from $C$ to $B{:}$
$$\begin{align*} A \to C \to B \end{align*}$$ - $A$ is an $M$-dimensional tensor with dimensions $m_1 \times m_2 \times \cdots \times m_M$
- $C$ is a $1$-dimensional tensor with $p$ elements
- $B$ is an $N$-dimensional tensor with dimensions $n_1 \times n_2 \times \cdots \times n_N$
- $m_1 m_2 \cdots m_M = p = n_1 n_2 \cdots n_N$
Using the zero-indexed row-major convention, we have
$$\begin{align*} A[i_1, i_2, \ldots, i_M] = C[i_1 (m_2 m_3 \cdots m_M) + i_2 (m_3 m_4 \cdots m_M) + \cdots + i_{M-1} (m_M) + i_M], \end{align*}$$ or, equivalently,
$$\begin{align*} C[k] &= A[i_1, i_2, \ldots, i_M] \\ i_1 &= \left\lfloor \dfrac{k}{m_2 m_3 \cdots m_M} \right\rfloor \\ i_2 &= \left\lfloor \dfrac{k - i_1 (m_2 m_3 \cdots m_M)}{m_3 m_4 \cdots m_M} \right\rfloor \\ i_3 &= \left\lfloor \dfrac{k - i_1 (m_2 m_3 \cdots m_M) - i_2 (m_3 m_4 \cdots m_M)}{m_4 m_5 \cdots m_M} \right\rfloor \\ &\vdots \end{align*}$$
So, to map from $A[i_1,i_2,\ldots,i_M]$ to $B[j_1,j_2,\ldots,j_N],$ we have the following formula:
$$\begin{align*} k &= i_1 (m_2 m_3 \cdots m_M) + i_2 (m_3 m_4 \cdots m_M) + \cdots + i_{M-1} (m_M) + i_M \\ j_1 &= \left\lfloor \dfrac{k}{n_2 n_3 \cdots n_N} \right\rfloor \\ j_2 &= \left\lfloor \dfrac{k - j_1 (n_2 n_3 \cdots n_N)}{n_3 n_4 \cdots n_N} \right\rfloor \\ j_3 &= \left\lfloor \dfrac{k - j_1 (n_2 n_3 \cdots n_N) - j_2 (n_3 n_4 \cdots n_N)}{n_4 n_5 \cdots n_N} \right\rfloor \\ &\vdots \end{align*}$$
Sanity Check: suppose
$$\begin{align*} A = \begin{bmatrix} 10 & 20 & 30 \\ 40 & 50 & 60 \end{bmatrix}, \quad B = \begin{bmatrix} 10 & 20 \\ 30 & 40 \\ 50 & 60 \end{bmatrix}. \end{align*}$$
Finding $B[j_1,j_2]$ corresponding to $A[1,0] = 40{:}$
$$\begin{align*} k &= 1(3) + 0 = 3 \\ j_2 &= \left\lfloor \dfrac{3}{2} \right\rfloor = 1 \\ j_1 &= 3 - 1(2) = 1 \\ B[1,1] &= 40 \quad \checkmark \end{align*}$$
Finding $B[j_1,j_2]$ corresponding to $A[0,2] = 30{:}$
$$\begin{align*} k &= 0(3) + 2 = 2 \\ j_2 &= \left\lfloor \dfrac{2}{2} \right\rfloor = 1 \\ j_1 &= 2 - 1(2) = 0 \\ B[1,0] &= 30 \quad \checkmark \end{align*}$$
Further Simplification: To simplify further, you can replace $k \to k - j_1 (n_2 n_3 \cdots n_N)$ with $k \to k \mod (n_2 n_3 \cdots n_N),$ and so on:
$$\begin{align*} k &= i_1 (m_2 m_3 \cdots m_M) + i_2 (m_3 m_4 \cdots m_M) + \cdots + i_{M-1} (m_M) + i_M \\ j_1 &= \left\lfloor \dfrac{k}{n_2 n_3 \cdots n_N} \right\rfloor \\ j_2 &= \left\lfloor \dfrac{k \mod (n_2 n_3 \cdots n_N)}{n_3 n_4 \cdots n_N} \right\rfloor \\ j_3 &= \left\lfloor \dfrac{k \mod (n_2 n_3 \cdots n_N) \mod (n_3 n_4 \cdots n_N)}{n_4 n_5 \cdots n_N} \right\rfloor \\ &\vdots \end{align*}$$