生成指定概率的0/1数组：元素为1的概率为1/(|i-j|+2)

问题描述 #

需要生成一个二维数组，其中每个位置(i,j)的元素为1的概率是1/(|i-j|+2)（i、j为行列索引），否则为0。但当前运行代码后输出的是随机大整数数组，不符合预期的仅含0和1的结果。

代码中的错误分析 #

1. np.empty的问题：np.empty((nr,nc), dtype=int)创建的是未初始化数组，内存中填充的是随机垃圾值，并非可直接填充的空数组。
2. np.append的误用：np.append不会修改原数组，而是返回新数组。循环中仅调用np.append却未将返回值赋值给values，导致原数组始终是初始化时的垃圾值。

修正后的代码 #

方式一：修复原循环逻辑 #

import numpy as np
import random

def seed_probabilities(i,j):
    rand_num = random.random()
    x = 1/(abs(i - j) + 2)
    return 1 if rand_num < x else 0

def create_world(nr, nc):
    # 用zeros初始化全0数组，避免垃圾值
    values = np.zeros((nr, nc), dtype=int)
    for i in range(nr):
        for j in range(nc):
            # 直接通过索引赋值，修改原数组
            values[i, j] = seed_probabilities(i,j)
    return values

方式二：更高效的Numpy广播实现（避免循环） #

如果处理大数组，推荐用Numpy向量化操作，比循环效率更高：

import numpy as np

def create_world(nr, nc):
    # 生成行索引矩阵和列索引矩阵
    i = np.arange(nr)[:, np.newaxis]
    j = np.arange(nc)
    # 计算每个位置的概率
    probabilities = 1/(np.abs(i - j) + 2)
    # 生成随机数矩阵并与概率比较，得到0-1数组
    return (np.random.rand(nr, nc) < probabilities).astype(int)

验证结果 #

调用create_world(6,7)后，输出会是类似这样的仅含0和1的数组：

array([[1, 0, 0, 1, 0, 0, 0],
       [0, 1, 0, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0, 1],
       [0, 1, 0, 1, 0, 0, 0],
       [1, 0, 0, 0, 1, 0, 0],
       [0, 0, 0, 0, 0, 1, 0]])

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