使用NumPy查找所有n維線條和對角線

Question

使用NumPy，我想列出長度爲k的n維數組的所有線條和對角線。

---

以長度爲3的以下三維數組爲例。

array([[[ 0, 1, 2], 
     [ 3, 4, 5], 
     [ 6, 7, 8]], 

     [[ 9, 10, 11], 
     [12, 13, 14], 
     [15, 16, 17]], 

     [[18, 19, 20], 
     [21, 22, 23], 
     [24, 25, 26]]]) 

對於這種情況，我想獲得所有以下類型的序列。對於任何特定情況，我希望獲得每種類型的所有可能序列。對於每種情況，所需序列的例子在下面的括號中給出。

• 1D線
  • x軸（0, 1, 2）
  • Y軸（0, 3, 6）
  • Z軸（0, 9, 18）
• 2D對角線
  • 的x/y軸（0, 4, 8，2, 4, 6）
  • X/Z軸（0, 10, 20，2, 10, 18）
  • Y/Z軸（0, 12, 24，6, 12, 18）
• 3D對角線
  • 的x/y/z軸（0, 13, 26， 2, 13, 24）

---

該解決方案應該被廣義化，以便它將生成一個數組的所有行和對角線，而不管數組的維數或長度（在所有維度上都是恆定的）。

Answer 1

這解決方案概括爲n

讓我們來重述一下這個pro瑕疵爲「找到索引列表」。

我們正在尋找一切形式的二維數組索引的

array[i[0], i[1], i[2], ..., i[n-1]] 

讓n = arr.ndim

哪裏i是形狀(n, k)

陣列，每個的i[j]可以是一個：

• 相同的索引重複n次，ri[j] = [j, ..., j]
• 的正向序列，fi = [0, 1, ..., k-1]
• 向後序列，bi = [k-1, ..., 1, 0]

隨着每個序列的形式^(ri)*(fi)(fi|bi|ri)*$的要求（使用正則表達式來概括它）。這是因爲：

• 必須有至少一個fi因此，「行」是不會重複選擇一個點
• 沒有biš來之前fi S，以避免得到扭轉線

---

def product_slices(n): 
    for i in range(n): 
     yield (
      np.index_exp[np.newaxis] * i + 
      np.index_exp[:] + 
      np.index_exp[np.newaxis] * (n - i - 1) 
     ) 

def get_lines(n, k): 
    """ 
    Returns: 
     index (tuple): an object suitable for advanced indexing to get all possible lines 
     mask (ndarray): a boolean mask to apply to the result of the above 
    """ 
    fi = np.arange(k) 
    bi = fi[::-1] 
    ri = fi[:,None].repeat(k, axis=1) 

    all_i = np.concatenate((fi[None], bi[None], ri), axis=0) 

    # inedx which look up every possible line, some of which are not valid 
    index = tuple(all_i[s] for s in product_slices(n)) 

    # We incrementally allow lines that start with some number of `ri`s, and an `fi` 
    # [0] here means we chose fi for that index 
    # [2:] here means we chose an ri for that index 
    mask = np.zeros((all_i.shape[0],)*n, dtype=np.bool) 
    sl = np.index_exp[0] 
    for i in range(n): 
     mask[sl] = True 
     sl = np.index_exp[2:] + sl 

    return index, mask 

---

應用到你的例子：

# construct your example array 
n = 3 
k = 3 
data = np.arange(k**n).reshape((k,)*n) 

# apply my index_creating function 
index, mask = get_lines(n, k) 

# apply the index to your array 
lines = data[index][mask] 
print(lines) 

array([[ 0, 13, 26], 
     [ 2, 13, 24], 
     [ 0, 12, 24], 
     [ 1, 13, 25], 
     [ 2, 14, 26], 
     [ 6, 13, 20], 
     [ 8, 13, 18], 
     [ 6, 12, 18], 
     [ 7, 13, 19], 
     [ 8, 14, 20], 
     [ 0, 10, 20], 
     [ 2, 10, 18], 
     [ 0, 9, 18], 
     [ 1, 10, 19], 
     [ 2, 11, 20], 
     [ 3, 13, 23], 
     [ 5, 13, 21], 
     [ 3, 12, 21], 
     [ 4, 13, 22], 
     [ 5, 14, 23], 
     [ 6, 16, 26], 
     [ 8, 16, 24], 
     [ 6, 15, 24], 
     [ 7, 16, 25], 
     [ 8, 17, 26], 
     [ 0, 4, 8], 
     [ 2, 4, 6], 
     [ 0, 3, 6], 
     [ 1, 4, 7], 
     [ 2, 5, 8], 
     [ 0, 1, 2], 
     [ 3, 4, 5], 
     [ 6, 7, 8], 
     [ 9, 13, 17], 
     [11, 13, 15], 
     [ 9, 12, 15], 
     [10, 13, 16], 
     [11, 14, 17], 
     [ 9, 10, 11], 
     [12, 13, 14], 
     [15, 16, 17], 
     [18, 22, 26], 
     [20, 22, 24], 
     [18, 21, 24], 
     [19, 22, 25], 
     [20, 23, 26], 
     [18, 19, 20], 
     [21, 22, 23], 
     [24, 25, 26]]) 

---

另一套很好的測試數據是np.moveaxis(np.indices((k,)*n), 0, -1)，這給出了一個數組，其中的每個值都是自己的索引

---

我之前解決了這個問題實施higher dimensional tic-tac-toe

Answer 2

In [1]: x=np.arange(27).reshape(3,3,3) 

選擇單個rows很簡單：

In [2]: x[0,0,:] 
Out[2]: array([0, 1, 2]) 
In [3]: x[0,:,0] 
Out[3]: array([0, 3, 6]) 
In [4]: x[:,0,0] 
Out[4]: array([ 0, 9, 18]) 

你可以在同一個索引列表尺寸迭代：

In [10]: idx=[slice(None),0,0] 
In [11]: x[idx] 
Out[11]: array([ 0, 9, 18]) 
In [12]: idx[2]+=1 
In [13]: x[idx] 
Out[13]: array([ 1, 10, 19]) 

看看在np.apply_along_axis代碼，看看它是如何實現這種的迭代。

重塑和拆分也可以生成rows的列表。對於某些方面，這可能需要一個transpose：

In [20]: np.split(x.reshape(x.shape[0],-1),9,axis=1) 
Out[20]: 
[array([[ 0], 
     [ 9], 
     [18]]), array([[ 1], 
     [10], 
     [19]]), array([[ 2], 
     [11], 
     ... 

np.diag可以從2D子陣

In [21]: np.diag(x[0,:,:]) 
Out[21]: array([0, 4, 8]) 
In [22]: np.diag(x[1,:,:]) 
Out[22]: array([ 9, 13, 17]) 
In [23]: np.diag? 
In [24]: np.diag(x[1,:,:],1) 
Out[24]: array([10, 14]) 
In [25]: np.diag(x[1,:,:],-1) 
Out[25]: array([12, 16]) 

獲得對角線探索np.diagonal直接應用到3D。直接爲陣列索引也很容易，用range和arange,x[0,range(3),range(3)]。

據我所知，沒有一個功能來逐步完成所有這些選擇。由於返回數組的維數可能不同，因此在編譯的numpy代碼中生成這樣的函數幾乎沒有意義。所以，即使有功能，它也會按照我所概述的方式逐步完成替代方案。

==============

所有1D線

x.reshape(-1,3) 
x.transpose(0,2,1).reshape(-1,3) 
x.transpose(1,2,0).reshape(-1,3) 

Y/Z的對角線和反對角線

In [154]: i=np.arange(3) 
In [155]: j=np.arange(2,-1,-1) 
In [156]: np.concatenate((x[:,i,i],x[:,i,j]),axis=1) 
Out[156]: 
array([[ 0, 4, 8, 2, 4, 6], 
     [ 9, 13, 17, 11, 13, 15], 
     [18, 22, 26, 20, 22, 24]]) 

Answer 3

np.einsum可用於構建所有這些ki表達式;例如：

# 3d diagonals 
print(np.einsum('iii->i', a)) 
# 2d diagonals 
print(np.einsum('iij->ij', a)) 
print(np.einsum('iji->ij', a))