有没有一种很好的方法可以在python中迭代几个列表(有时它们是单值的)

时间:2013-03-01 22:03:41

标签: python numpy

所以基本上我的问题与'zip'(或izip)有关,而这个问题在之前被提出过....

Is there a better way to iterate over two lists, getting one element from each list for each iteration?

如果我有两个变量 - 它们是长度为n的1d数组,或者是单个值,我如何循环它们以便我返回n个值。

'zip'有点做我想要的 - 除了当我传入一个值,并且它抱怨一个数组。

我有一个我正在瞄准的例子 - 基本上我有一个c函数,它比python做更有效的计算。我希望它像一些numpy函数一样 - 用数组和标量的混合物来处理,所以我为它编写了一个python包装器。但是 - 就像我说'拉链'失败了。我想原则上我可以对输入s进行一些测试,并为标量和数组的每个变体写一个不同的语句 - 但看起来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.braodcast - 因为从我的角度看它似乎是最简单的......

3 个答案:

答案 0 :(得分:4)

由于您使用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)]

答案 1 :(得分:1)

如果您想强制广播,可以使用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循环中运行。

答案 2 :(得分:0)

将每个参数转换为序列的装饰器可能会有所帮助。这是普通的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)])