如何高效更新多维NumPy数组的多份拷贝并实现指定位置随机增量

核心目标 #

给定多维数组与对应同尺寸的0-1索引集，将数组每一行复制n次后，仅对索引值为1的元素执行随机增量操作，增量的取值范围是1到upper_bound - 原元素值，保证增量后的值不超过upper_bound。
示例：数组为[[3,6,7,8], [1,32,45,7]]，索引集为[[1,0,1,1], [0,0,1,1]]，仅索引为1的位置会增加随机值。

现有实现 #

import time
import random
import numpy as np

def foo(arr, upper_bound, index_set, first_set_size, sec_set_size, limit):
    iter =0

    my_array = np.zeros((first_set_size*sec_set_size, limit)) #每行被复制sec_set_size次
    it =0
    for i in range(first_set_size):
        for j in range(sec_set_size):
            my_array[it] = arr[i] #复制对应行的元素
            for k in range(limit):
                if index_set[i][k]==1: #更新索引为1的元素
                    temp = arr[i][k]   #获取当前值
                    my_array[it][k]  =temp + random.randint(1,upper_bound-temp) #实际使用fastrand.pcg32bounded生成随机数
            it +=1
    return my_array


upper_bound = 50
limit = 1000
first_set_size= 100
sec_set_size = 50
arr = np.random.randint(25, size=(first_set_size, limit)) #生成整型数组
index_set= np.array([[random.randint(0,1) for j in range(limit)] for i in range(first_set_size)]) #每个元素对应0或1的索引

start_time = time.time() #统计函数耗时
result = foo(arr, upper_bound,index_set, first_set_size, sec_set_size, limit)
print("time taken: %s " % (time.time() - start_time))

该实现采用三层Python循环，当limit和集合规模较大时，运行耗时可达数分钟，性能不足。

示例验证 #

初始输入数组 #

[[11 23 24 17  0]
 [ 1 23 12 19  5]
 [20 15  1 17 17]
 [ 3  8  7  0 24]]

对应0-1索引集 #

[[1 0 0 0 1]
 [1 0 1 0 0]
 [1 1 1 1 0]
 [0 1 0 1 1]]

输出结果示例（sec_set_size=5） #

[[39. 23. 24. 17. 44.]
 [50. 23. 24. 17. 27.]
 [42. 23. 24. 17. 24.]
 [45. 23. 24. 17. 11.]
 [49. 23. 24. 17. 43.]
 [23. 23. 44. 19.  5.]
 [10. 23. 37. 19.  5.]
 [14. 23. 29. 19.  5.]
 [12. 23. 22. 19.  5.]
 [ 5. 23. 15. 19.  5.]
 [36. 45. 26. 37. 17.]
 [24. 40. 35. 38. 17.]
 [34. 20. 24. 31. 17.]
 [27. 16.  9. 20. 17.]
 [37. 37.  6. 37. 17.]
 [ 3. 50.  7. 46. 47.]
 [ 3. 13.  7. 37. 44.]
 [ 3. 23.  7. 32. 29.]
 [ 3. 10.  7. 22. 41.]
 [ 3. 22.  7. 32. 41.]]

核心思路是替换所有Python层循环为numpy向量化操作，利用广播特性批量处理数据：

import numpy as np
import time

def fast_foo(arr, upper_bound, index_set, sec_set_size):
    # 1. 按行重复sec_set_size次，得到基础扩展数组
    expanded_arr = np.repeat(arr, sec_set_size, axis=0)
    # 2. 索引集同样按行重复，和扩展数组形状对齐
    expanded_index = np.repeat(index_set, sec_set_size, axis=0)
    # 3. 生成所有位置的最大允许增量：upper_bound - 原值，仅索引为1的位置有效
    max_delta = upper_bound - expanded_arr
    # 4. 生成随机增量，范围[1, max_delta]，不需要修改的位置生成0即可
    rng = np.random.default_rng()
    delta = rng.integers(low=1, high=max_delta+1, size=expanded_arr.shape) * expanded_index
    # 5. 加回原数组得到结果
    return expanded_arr + delta

# 测试参数和原代码一致
upper_bound = 50
limit = 1000
first_set_size= 100
sec_set_size = 50
arr = np.random.randint(25, size=(first_set_size, limit))
index_set= np.random.randint(0, 2, size=(first_set_size, limit))

start_time = time.time()
result = fast_foo(arr, upper_bound, index_set, sec_set_size)
print("time taken: %s " % (time.time() - start_time))

性能对比 #

相同测试参数下，原实现耗时约2.7秒，优化后的实现耗时仅约0.01秒，性能提升270倍以上；规模越大性能优势越明显，完全避免了数分钟的等待。
如果需要使用fastrand.pcg32bounded进一步提升随机数生成速度，可以自行替换随机增量生成的部分，逻辑保持一致即可。

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