具有广播维度灵活性的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 循环，我只能考虑为if 语句中x、y 和z 的每个组合编写一个for 循环.这可能会使更复杂的数学表达式中的更多变量变得混乱。

有没有比使用许多 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