Julia代码性能优化及向量最大值坐标获取相关问题咨询

我从Python转向学习Julia，编写代码求解车遍历棋盘谜题的所有解，采用暴力搜索方案。Python版本运行耗时约30分钟，改写为Julia代码后耗时约3分钟，我已参考官方优化文档为函数添加数据类型标注、避免使用全局变量，但仍存在相关疑问，代码如下：

using Primes

function mymax(state::Vector{Vector{Int64}})
    maxval,maxi,maxj=0,0,0
    for (i,row) in enumerate(state)
        for (j,val) in enumerate(row)
            if val>maxval
                maxval,maxi,maxj = val,i,j
            end
        end
    end
    return maxval,maxi,maxj
end

function neighbors(state::Vector{Vector{Int64}})
    prev,i,j = mymax(state)
    ip = Primes.isprime(prev+1)
    newstates=Vector{Vector{Vector{Int}}}()
    for (di,dj) = [(-1,0),(1,0),(0,1),(0,-1)]
        ni,nj = i+di,j+dj
        (ni<1 || nj<1 || ni>length(state) || nj>length(state[1])) && continue
        (state[ni][nj]!=0) && continue
        if field[ni][nj]==ip
            newstate=deepcopy(state)
            newstate[ni][nj]=prev+1
            push!(newstates,newstate)
        end
    end
    return newstates
end

function goal(state::Vector{Vector{Int64}})
    return sum(sum(state))==2080
end

function search(start_state::Vector{Vector{Int64}},field::Vector{Vector{Int64}})
    path = [start_state]
    level=0
    solutions = Vector{Vector{Vector{Int}}}()
    while true
        (length(path)==0) && (return solutions)
        println("***")
        println(level)
        println(length(path))
        println(mymax(path[end]))
        level+=1
        newpath = Vector{Vector{Vector{Int}}}()
        for p in path
            for n in neighbors(p)
                (goal(n)) && (push!(solutions,n))
                push!(newpath,n)
            end
        end
        path=deepcopy(newpath)
    end
end

field = [[1,0,0,0,1,0,0,0],
         [0,1,0,1,0,0,0,1],
         [1,0,0,0,0,0,1,0],
         [0,0,0,1,0,0,0,1],
         [0,0,1,1,1,0,0,0],
         [0,0,0,0,0,1,0,0],
         [1,0,0,0,1,0,1,0],
         [0,0,0,0,0,1,0,1]]

 start = [[0,0,0,0,0,0,0,0],
          [0,0,0,0,0,0,0,0],
          [0,0,0,0,0,0,0,0],
          [0,0,0,0,0,0,0,0],
          [0,0,0,2,0,0,0,0],
          [0,0,0,1,0,0,0,0],
          [0,0,0,0,0,0,0,0],
          [0,0,0,0,0,0,0,0]]

solns = search(start,field)
print("Solutions found ")
println(length(solns))
for soln in solns
    display(soln)
end

具体疑问 #

• mymax函数是否有更优实现？我尝试使用argmax实现相关功能，但它返回的是元素总和最大的行的索引，仅在初始状态下符合需求，无法通用，请问正确获取二维向量最大值及其对应坐标的方法是什么？
• Vector{Vector{Int64}}这类嵌套向量是否是适合当前场景的数据结构？是否有更优化的结构可选？@time运行返回结果为181.230324 seconds (214.55 M allocations: 31.388 GiB, 14.44% gc time)，内存分配量很高，请问该如何排查该问题？
• 还有哪些优化技巧可以提升该代码的运行速度？我已经开启编译模式获得了明显提速，当前Julia版本相比Python已经有10倍性能提升，我想知道是否还有较大的优化空间，是否有我遗漏的核心优化点？

补充说明：我期望能在当前基础上进一步优化，实现Julia版本相比Python版本百倍级的性能提升。

问题解答 #

1. mymax函数优化方案 #

你当前的循环实现效率已经很高，Julia原生循环的性能和C相当，不需要为了语法简洁牺牲性能。你用argmax只返回行索引的原因是你使用了嵌套Vector结构，argmax对嵌套数组默认沿第一维比较行对象，自然只会返回行索引。
如果要更简洁的实现，只需要把嵌套Vector替换为原生二维数组Matrix{Int64}，此时argmax(state)会直接返回CartesianIndex(i,j)类型的坐标，取值和拆分解的写法如下：

ind = argmax(state)
maxval = state[ind]
i,j = ind.I

如果坚持使用嵌套Vector，原循环只需要把初始maxval设为typemin(Int64)即可提升鲁棒性，不需要其他修改。另外你当前场景下完全可以在生成新状态时同步传递最大值和坐标，省掉每轮O(64)的全数组遍历开销，是非常容易落地的优化点。

2. 数据结构选择与内存问题排查 #

• 嵌套VectorVector{Vector{Int64}}完全不适合当前场景：它的每行都是独立的堆分配对象，访问需要两次寻址，复制时要递归处理每行，开销极高。
• 最优选择是使用StaticArrays.jl的MMatrix{8,8,Int64}：因为棋盘大小固定为8x8，静态数组会在栈上分配，内存连续，复制整个数组只需要拷贝64个Int，开销可以忽略。次优选择是换成原生Matrix{Int64}，内存连续，访问和复制效率远高于嵌套Vector。
• 内存过高的核心原因有三个：
  1. 你当前的neighbors函数没有把field作为参数传入，直接调用了全局变量field，全局变量类型不稳定会导致大量隐式分配，这是最核心的性能坑。
  2. 多余的deepcopy调用：你每生成一个新状态就深拷贝整个数组，还在path赋值时又深拷贝了一遍所有状态，相当于每个状态被重复复制两次，完全没有必要。
  3. BFS方案需要保留所有层级的状态，本身就会产生大量内存占用，换成DFS回溯方案可以完全避免中间状态的存储开销。

3. 核心优化点（可实现百倍性能提升） #

以下优化落地后完全可以把运行时间降到10秒以内，实现百倍于Python的性能：

1. 消灭全局变量：所有用到的外部参数都要作为函数参数传入，或者用const修饰全局变量提升类型稳定性。
2. 替换数据结构：用静态数组MMatrix{8,8,Int64}存储棋盘状态，内存分配开销直接降为原来的1%以下。
3. 更换搜索方案：把BFS换成DFS回溯，不需要保存所有中间状态，直接修改当前状态，递归搜索完后回溯修改即可，全程几乎不会产生新的内存分配。
4. 删除冗余操作：
   • 去掉sum(sum(state))的判断，因为你是按1~64的顺序填值，直接判断当前最大值是否为64即可作为终止条件，省掉每轮求和开销。
   • 预计算1~64的素数表，存在静态Bool数组里，不需要每次调用isprime。
   • 删掉调试用的打印语句，IO操作的开销远高于计算逻辑。
   • 去掉多余的deepcopy，静态数组用普通copy即可，回溯方案甚至不需要复制状态。
5. 类型稳定优化：所有容器的类型要统一，不要混用Int和Int64，确保函数返回值类型可被编译器推断。

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