如何在Julia中向量化多输出函数并调整输出格式？

问题背景 #

你定义了一个返回多个值的函数：

function z(a,b)
    x_1 = a*b
    x_2 = a + b
    x_3 = a^b
    return x_1, x_2, x_3
end

当对输入A=2和B=[1,4]使用广播调用z.(A,B)时，得到的结果是[(2,3,2),(8,6,16)]——这是每个输入对对应的结果元组组成的数组。但你需要的是每个输出分量单独聚合：x_1=(2,8), x_2=(3,6), x_3=(2,16)。

解决方案 #

方法1：用zip解包广播结果 #

这是最简洁快捷的方式，利用zip将广播得到的元组数组按分量拆分：

A = 2
B = [1,4]

x1, x2, x3 = zip(z.(A,B)...)

执行后x1就是(2,8)，x2是(3,6)，x3是(2,16)，完全符合需求。如果需要数组而非元组，只需用collect转换：

x1_arr, x2_arr, x3_arr = collect.(zip(z.(A,B)...))

方法2：自定义高效向量化函数 #

如果处理大规模数据，想要避免中间元组数组的开销，可以直接写循环版本的向量化函数，预分配结果数组：

function z_vec(a, b_vec)
    # 预分配与输入类型匹配的结果数组
    x1 = similar(b_vec, promote_type(typeof(a), eltype(b_vec)))
    x2 = similar(x1)
    x3 = similar(x1)
    
    for i in eachindex(b_vec)
        x1[i], x2[i], x3[i] = z(a, b_vec[i])
    end
    
    return x1, x2, x3
end

# 调用示例
x1, x2, x3 = z_vec(A, B)

这种方式性能最优，直接操作数组元素，无额外元组分配开销。

方法3：拆分分量单独广播 #

如果每个输出分量逻辑独立，也可以直接对每个运算单独广播：

x1 = A .* B
x2 = A .+ B
x3 = A .^ B

这种方式代码直观，但原函数逻辑复杂时会导致代码重复，仅适合简单场景。

内容的提问来源于stack exchange，提问作者ValientProcess