高次元配列のナンバリング

次元毎に要素数の異なる配列のある要素に整数でインデックスを振りたいという状況がある
- 強化学習のソフトマックス行動選択を書いていたらそういう状況が出てきた

かなり悩んでから次のようなアルゴリズムを書いた
$K$ 次元配列の $k$ 次元の要素数を $N_k$ とすると，各次元のインデックス $(i_1,\cdots,i_k,\cdots,i_K)$ に対するナンバリングは

${\rm Num}(i_1,\cdots,i_k,\cdots,i_K) = i_K + \sum_{k=1}^{K-1} i_k \prod_{l=k}^{K-1} N_l$

LISPでの実装

(defun dimension-list->index (dimension-list index-list)
  (nlet itr ((dim-lst (cdr dimension-list))
	     (i-lst index-list)
	     (product 0))
    (if (null dim-lst)
	(+ product (car i-lst))
	(itr (cdr dim-lst) (cdr i-lst)
	     (+ (apply #'* (car i-lst) dim-lst)
		product)))))

これの逆写像を考える
- 順写像はかけて → 足す
- 逆写像は余りをとってひく → 割る

(defun scalar-index->index-list (dimension-list scalar-index)
  (nlet itr ((dim-lst (reverse dimension-list))
	     (i-lst '())
	     (num scalar-index))
    (if (= (length dim-lst) 1)
	(cons num i-lst)
	(let ((remainder (mod num (car dim-lst))))
	  (itr (cdr dim-lst)
	       (cons remainder i-lst)
	       (/ (- num remainder) (car dim-lst)))))))

例：3×4×2次元配列のナンバリング

CL-USER> (sfor (i 0 2)
	   (sfor (j 0 3)
	     (sfor (k 0 1)
	       (let* ((sindex (index-list->scalar-index '(3 4 2) (list i j k)))
		      (index-lst (scalar-index->index-list '(3 4 2) sindex)))
		 (format t "(~A ~A ~A) => ~A => ~A~%" i j k sindex index-lst)))))

(0 0 0) => 0 => (0 0 0)
(0 0 1) => 1 => (0 0 1)
(0 1 0) => 2 => (0 1 0)
(0 1 1) => 3 => (0 1 1)
(0 2 0) => 4 => (0 2 0)
(0 2 1) => 5 => (0 2 1)
(0 3 0) => 6 => (0 3 0)
(0 3 1) => 7 => (0 3 1)
(1 0 0) => 8 => (1 0 0)
(1 0 1) => 9 => (1 0 1)
(1 1 0) => 10 => (1 1 0)
(1 1 1) => 11 => (1 1 1)
(1 2 0) => 12 => (1 2 0)
(1 2 1) => 13 => (1 2 1)
(1 3 0) => 14 => (1 3 0)
(1 3 1) => 15 => (1 3 1)
(2 0 0) => 16 => (2 0 0)
(2 0 1) => 17 => (2 0 1)
(2 1 0) => 18 => (2 1 0)
(2 1 1) => 19 => (2 1 1)
(2 2 0) => 20 => (2 2 0)
(2 2 1) => 21 => (2 2 1)
(2 3 0) => 22 => (2 3 0)
(2 3 1) => 23 => (2 3 1)
NIL