具有广播尺寸灵活性的For循环

Question

与广播的辅助下,下面的代码将工作是否x,y和z是标量,大小的矢量n,或它们的任何组合.

b = zeros(n)
b .= x.*y.*z .+ x

但是,我想要一个for循环.以下for循环仅在x大小为矢量时才有效n,y是标量,并且z是标量.

for i = 1:n
    b[i] = x[i]*y*z + x[i]
end

要写入相当于b .= x.*y.*z .+ x为一个for循环的任何情况下,我只能想到写一个for循环的每一个组合的x,y和z内if语句.在更复杂的数学表达式中,这可能会使更多变量变得混乱.

有没有比使用许多if语句更优雅的方式来做我喜欢的事情？

Answer 1

您可以定义一个包装器类型,如果包装变量是数组,则索引到它将给出数组索引,并为标量的所有索引重复相同的值.我在下面有一个例子,但它可能不如使用效率高broadcast.并且它不检查数组长度是否一致.但是,自定义包装类型可以缓解这种情况.

julia> function f(x,y,z)
           lx,ly,lz = length(x),length(y),length(z)
           maxlen = max(lx,ly,lz)
           cx = cycle(x)
           cy = cycle(y)
           cz = cycle(z)
           b = zeros(maxlen)
           @inbounds for (xi,yi,zi,i) in zip(cx,cy,cz,1:maxlen)
               b[i] = xi*yi*zi+xi
           end
           return b
       end
f (generic function with 1 method)

julia> f(1:3,21,2)
3-element Array{Float64,1}:
  43.0
  86.0
 129.0