在Python中，有没有好的方法遍历多个列表（有时它们是单一值）？
我的问题主要和'zip'（或者说izip）有关，还有之前有人问过的这个问题……
有没有更好的方法来遍历两个列表，每次从每个列表中取一个元素？
如果我有两个变量，它们要么是长度为n的一维数组，要么是一个单独的值，我该如何循环遍历它们，以便返回n个值呢？
'zip'在某种程度上可以满足我的需求——但当我传入一个单独的值和一个数组时，它就会出错。
下面是我想要的一个例子——基本上我有一个C语言的函数，它的计算效率比Python高。我希望它能像一些numpy函数那样，能够处理数组和单个值的混合，所以我为它写了一个Python的包装函数。不过，正如我所说的，'zip'不太适用。我想原则上我可以对输入进行一些测试，然后为每种标量和数组的组合写不同的语句，但我觉得Python应该有更聪明的解决办法…… ;) 有什么建议吗？
"""
Example of zip problems.
"""
import numpy as np
import time
def cfun(a, b) :
"""
Pretending to be c function which doesn't deal with arrays
"""
if not np.isscalar(a) or not np.isscalar(b) :
raise Exception('c is freaking out')
else :
return a+b
def pyfun(a, b) :
"""
Python Wrappper - to deal with arrays input
"""
if not np.isscalar(a) or not np.isscalar(b) :
return np.array([cfun(a_i,b_i) for a_i, b_i in zip(a,b)])
else :
return cfun(a, b)
return cfun(a,b)
a = np.array([1,2])
b= np.array([1,2])
print pyfun(a, b)
a = [1,2]
b = 1
print pyfun(a, b)
编辑：
非常感谢大家的建议。我觉得我应该使用np.broadcast来解决这个问题——因为从我的角度来看，这似乎是最简单的办法……
3 个回答
一个装饰器可以选择性地把每个参数转换成一个序列，这样可能会很有帮助。下面是普通的Python版本（不是numpy的）：
# TESTED
def listify(f):
def dolistify(*args):
from collections import Iterable
return f(*(a if isinstance(a, Iterable) else (a,) for a in args))
return dolistify
@listify
def foo(a,b):
print a, b
foo( (1,2), (3,4) )
foo( 1, [3,4] )
foo( 1, 2 )
所以，在你的例子中，我们需要用 not np.isscalar 作为条件，用 np.array 作为修改器。因为有了这个装饰器，pyfun 总是能接收到一个数组。
#UNTESTED
def listify(f):
def dolistify(*args):
from collections import Iterable
return f(*(np.array([a]) if np.isscalar(a) else a for a in args))
return dolistify
@listify
def pyfun(a, b) :
"""
Python Wrappper - to deal with arrays input
"""
return np.array([cfun(a_i,b_i) for a_i, b_i in zip(a,b)])
或者你也可以把同样的想法应用到 zip 上：
#UNTESTED
def MyZip(*args):
return zip(np.array([a]) if np.isscalar(a) else a for a in args)
def pyfun(a, b) :
"""
Python Wrappper - to deal with arrays input
"""
return np.array([cfun(a_i,b_i) for a_i, b_i in MyZip(a,b)])
如果你想强制进行广播操作，可以使用 numpy.lib.stride_tricks.broadcast_arrays。你可以继续使用你的 cfun：
def pyfun(a, b) :
if not (np.isscalar(a) and np.isscalar(b)) :
a_bcast, b_bcast = np.lib.stride_tricks.broadcast_arrays(a, b)
return np.array([cfun(j, k) for j, k in zip(a_bcast, b_bcast)])
return cfun(a, b)
现在：
>>> pyfun(5, 6)
11
>>> pyfun(5, [6, 7, 8])
array([11, 12, 13])
>>> pyfun([3, 4, 5], [6, 7, 8])
array([ 9, 11, 13])
对于你的具体应用，可能没有比Rob的纯Python方法更好的选择，因为你的函数仍然是在Python循环中运行。
因为你使用了numpy，所以在遍历多个数组和标量时，不需要用到zip()这个函数。你可以直接使用numpy.broadcast()来处理。
In [5]:
list(np.broadcast([1,2,3], 10))
Out[5]:
[(1, 10), (2, 10), (3, 10)]
In [6]:
list(np.broadcast([1,2,3], [10, 20, 30]))
Out[6]:
[(1, 10), (2, 20), (3, 30)]
In [8]:
list(np.broadcast([1,2,3], 100, [10, 20, 30]))
Out[8]:
[(1, 100, 10), (2, 100, 20), (3, 100, 30)]