Python:计算不同长度列表列表中第n个元素的平均值

Question

假设我有以下列表:

a = [ 
      [1, 2, 3],
      [2, 3, 4],
      [3, 4, 5, 6] 
    ]

我想要数组中每个第 n 个元素的平均值。然而，当想以简单的方式做到这一点时，由于长度不同，Python 产生了越界错误。我通过为每个数组提供最长数组的长度并用 None 填充缺失值来解决这个问题。

不幸的是，这样做无法计算平均值，因此我将数组转换为屏蔽数组。下面显示的代码有效，但看起来相当麻烦。

import numpy as np
import numpy.ma as ma

a = [ [1, 2, 3],
      [2, 3, 4],
      [3, 4, 5, 6] ]

# Determine the length of the longest list
lenlist = []
for i in a:
    lenlist.append(len(i))
max = np.amax(lenlist)

# Fill each list up with None's until required length is reached
for i in a:
    if len(i) <= max:
        for j in range(max - len(i)):
            i.append(None)

# Fill temp_array up with the n-th element
# and add it to temp_array 
temp_list = []
masked_arrays = []
for j in range(max):
    for i in range(len(a)):
        temp_list.append(a[i][j])
    masked_arrays.append(ma.masked_values(temp_list, None))
    del temp_list[:]

# Compute the average of each array 
avg_array = []
for i in masked_arrays:
    avg_array.append(np.ma.average(i))

print avg_array

有没有办法更快地做到这一点？列表的最终列表将包含 600000 个“行”和最多 100 个“列”，因此效率非常重要:-)。

Answer 1

tertools.izip_longest会为你做所有的填充，这样你的代码就可以减少到:

import numpy as np
import numpy.ma as ma
from itertools import izip_longest

a = [ [1, 2, 3],
      [2, 3, 4],
      [3, 4, 5, 6] ]


averages = [np.ma.average(ma.masked_values(temp_list, None)) for temp_list in izip_longest(*a)]

print(averages)
[2.0, 3.0, 4.0, 6.0]

不知道关于 numpy 逻辑的最快方法是什么，但这肯定会比您自己的代码高效得多。

如果你想要一个更快的纯 python 解决方案:

from itertools import izip_longest, imap

a = [[1, 2, 3],
     [2, 3, 4],
     [3, 4, 5, 6]]


def avg(x):
    x = filter(None, x)
    return sum(x, 0.0) / len(x)


filt = imap(avg, izip_longest(*a))

print(list(filt))
[2.0, 3.0, 4.0, 6.0]

如果数组中的 0 不起作用，因为 0 将被视为 Falsey，在这种情况下，您将不得不使用 list comp 进行过滤，但它仍然会更快:

def avg(x):
    x = [i for i in x if i is not None]
    return sum(x, 0.0) / len(x)

filt = imap(avg, izip_longest(*a))