这两天尝试将一个MATLAB工程转成Python，卡在一个bug上很久，最后发现是二者在矩阵索引上存在区别。这里我采用的是Python的NumPy库。

我们经常会通过一定的逻辑关系索引矩阵中的元素，并获取这些元素的位置信息，即indices。例如，

1
2
3
4
5
6
% MATLAB实现
x = resize(0:8, 3, 3);
% 一维索引
idx = find(x >= 3);
% 二维索引
[idr, idc] = find(x >= 3);

对应的python实现为，

1
2
3
4
5
import numpy as np
x = np.arange(9).reshape(3,3)
# 只有二维索引
[idr, idc] = np.where(x >= 3)

这里可以看出，MATLAB提供一种一维索引，即将二维矩阵以列为先，行次之转为一维向量，输出的索引对应该一维向量中元素所在位置。（这种处理方法经常可以用于加速运算，缩短运行时间。）然而，NumPy的where方法根据矩阵的维数提供对应axis的索引，没有MATLAB这种一维索引的输出。

因此，在python中使用np.where进行矩阵元素索引时，要注意如下两点，

  1. np.where的输出是一个列表;
  2. np.where的输出列表的元素个数与矩阵的维数对应。

Reference

[1] numpy.where