NumPy - 遍历 2D 列表并打印（行、列）索引答案

【问题标题】：NumPy - Iterate over 2D list and print (row,column) indexNumPy - 遍历 2D 列表并打印（行、列）索引
【发布时间】：2017-04-26 15:09:54
【问题描述】：

我在使用NumPy 和/或Pandas 处理2D 列表时遇到困难：

获取所有元素的唯一组合sum，无需再次从同一行中选择（下面的数组应该是81个组合）。
打印组合中每个元素的行和列。

例如：

arr = [[1, 2, 4], [10, 3, 8], [16, 12, 13], [14, 4, 20]]

(1,3,12,20), Sum = 36 and (row, col) =  [(0,0),(1,1),(2,1),(3,2)]

(4,10,16,20), Sum = 50 and (row, col) =[(0,2),(1,0),(2,0),(3,2)]

【问题讨论】：

我尝试了常规的 python for 循环。但我需要使用 numpy/panda 的高性能数据结构。我搜索了很多论坛，但找不到如何遍历所有元素组合。

标签： python pandas numpy

【解决方案1】：

你可以使用itertools中的product函数：

from itertools import product    
y = [sum(p) for p in product(*arr)]

len(y)
# 81

列表较小的示例：

arr = [[1,2],[3,4],[5,6]]
[sum(p) for p in product(*arr)]
# [9, 10, 10, 11, 10, 11, 11, 12]

【讨论】：

谢谢，它不会打印有助于得出该总和的元素的 (row,col) 索引。我该怎么做呢？

【解决方案2】：

通过创建所有此类组合并求和的方法：这是使用 itertools.product 和 array-indexing 的矢量化方法 -

from itertools import product

a = np.asarray(arr)  # Convert to array for ease of use and indexing
m,n = a.shape
combs = np.array(list(product(range(n), repeat=m)))
out = a[np.arange(m)[:,None],combs.T].sum(0)

示例运行 -

In [296]: arr = [[1, 2, 4], [10, 3, 8], [16, 12, 13], [14, 4, 20]]

In [297]: a = np.asarray(arr)
     ...: m,n = a.shape
     ...: combs = np.array(list(product(range(n), repeat=m)))
     ...: out = a[np.arange(m)[:,None],combs.T].sum(0)
     ...: 

In [298]: out
Out[298]: 
array([41, 31, 47, 37, 27, 43, 38, 28, 44, 34, 24, 40, 30, 20, 36, 31, 21,
       37, 39, 29, 45, 35, 25, 41, 36, 26, 42, 42, 32, 48, 38, 28, 44, 39,
       29, 45, 35, 25, 41, 31, 21, 37, 32, 22, 38, 40, 30, 46, 36, 26, 42,
       37, 27, 43, 44, 34, 50, 40, 30, 46, 41, 31, 47, 37, 27, 43, 33, 23,
       39, 34, 24, 40, 42, 32, 48, 38, 28, 44, 39, 29, 45])

节省内存的方法：这是一种无需创建所有这些组合而是使用即时broadcasted 求和的方法，其理念深受this other post 的启发-

a = np.asarray(arr)
m,n = a.shape
out = a[0]
for i in range(1,m):
    out = out[...,None]  + a[i]
out.shape = out.size # Flatten

【讨论】：