Fortran一维数组按条件动态扩容的高效替代方案咨询

你遇到的问题本质是：mat=[mat,i]的数组追加方式会频繁触发内存重分配与数据复制，时间复杂度为O(k²)（k为最终元素数），大数据量下必然缓慢；而文件中转法虽然快，但会产生数百GB的磁盘文件，显然不可行。以下是两种高效且无额外磁盘开销的解决方案：

方案一：预统计元素数量，一次性分配数组 #

思路 #

先遍历一次数据集，统计符合条件的元素总数，再直接分配对应大小的数组，最后二次遍历填充元素。这种方式仅需一次内存分配，时间复杂度为O(n)，完全避免磁盘IO。

示例代码（与原程序结果一致） #

program test
    implicit none

    integer, parameter :: ndig=8
    integer, parameter :: ilong=selected_int_kind(ndig)
    integer (ilong), allocatable :: mat(:), rand_vals(:)
    integer (ilong), parameter :: n=2000000_ilong
    integer (ilong) :: i, cn
    real(8) :: START_CLOCK, STOP_CLOCK

    call cpu_time(START_CLOCK)

    ! 预先生成所有随机数，保证两次遍历的选择结果一致
    allocate(rand_vals(n))
    call random_number(rand_vals)

    ! 第一步：统计符合条件的元素数量
    cn = count(rand_vals < 0.5d0)

    ! 第二步：分配数组并填充
    allocate(mat(cn))
    cn = 0
    do i=1,n
        if (rand_vals(i) < 0.5d0) then
            cn = cn + 1
            mat(cn) = i
        end if
    end do

    deallocate(rand_vals)

    call cpu_time(STOP_CLOCK)
    print *, 'run took:', (STOP_CLOCK - START_CLOCK)/60.0d0, 'minutes.'

end program test

方案二：动态扩容数组（适合无法预统计的场景） #

思路 #

如果无法提前统计符合条件的元素数量（比如条件依赖实时计算结果），可以先预分配一个初始容量的数组，当元素数超过容量时，按固定比例扩容（比如1.5倍或2倍），减少内存重分配的次数。这种方式的时间复杂度接近O(n)，仅需少量额外内存。

示例代码 #

program test
    implicit none

    integer, parameter :: ndig=8
    integer, parameter :: ilong=selected_int_kind(ndig)
    integer (ilong), allocatable :: mat(:)
    integer (ilong), parameter :: n=2000000_ilong
    integer (ilong) :: i, cn, capacity
    real(8) :: z, START_CLOCK, STOP_CLOCK

    call cpu_time(START_CLOCK)

    ! 初始预分配容量（可根据业务场景调整，这里按最大可能值的1.2倍估算）
    capacity = n / 2 * 12 / 10
    allocate(mat(capacity))
    cn = 0

    do i=1,n
        call random_number(z)
        if (z < 0.5d0) then
            cn = cn + 1
            ! 元素数超过容量时，按1.5倍扩容
            if (cn > capacity) then
                capacity = int(capacity * 1.5d0, ilong)
                block
                    integer(ilong), allocatable :: temp(:)
                    allocate(temp(capacity))
                    temp(1:cn-1) = mat(1:cn-1)
                    call move_alloc(temp, mat)  ! 转移内存所有权，避免额外复制
                end block
            end if
            mat(cn) = i
        end if
    end do

    ! 裁剪数组到实际元素数量，释放多余内存
    if (cn < capacity) then
        block
            integer(ilong), allocatable :: temp(:)
            allocate(temp(cn))
            temp = mat(1:cn)
            call move_alloc(temp, mat)
        end block
    end if

    call cpu_time(STOP_CLOCK)
    print *, 'run took:', (STOP_CLOCK - START_CLOCK)/60.0d0, 'minutes.'

end program test

方案对比 #

• 方案一：性能最优，内存开销最小，代码简洁，适合可以提前统计元素数量的场景。
• 方案二：灵活性高，适合无法预统计的场景，扩容策略可根据实际情况调整（比如2倍扩容会减少分配次数，但会占用更多临时内存）。
• 文件中转法：仅适合内存完全无法容纳最终数组的极端场景，否则不推荐使用。

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