Python实现Kd×Kd矩阵对角线K×K块置零的代码咨询

问题背景 #

原有代码针对d×d矩阵生成约束，强制对角线元素为0，其余元素非负：

bnds = [(0, 0) if i == j else (0, None) for _ in range(2) for i in range(d) for j in range(d)]

现需修改为处理Kd×Kd的分块矩阵：矩阵由d个K×K块组成，要求所有对角位置的K×K块为零矩阵（比如d=20、K=5时，第1-5行1-5列、6-10行6-10列等块需全为0），其余元素非负。

你的尝试代码 #

bnds = []
for i in range(1, K*d+1):
    for j in range(1, K*d+1):
        block_row = (i - 1) // K
        block_col = (j - 1) // K
        row_start = block_row * K + 1
        row_end = row_start + K
        col_start = block_col * K + 1
        col_end = col_start + K
        if block_row == block_col:
            bnds.append((0, 0)) 
        else:
            bnds.append((0, None))

代码点评与优化 #

1. 核心逻辑正确：通过(i-1)//K和(j-1)//K计算元素所属的块行、块列，当块行与块列相等时，判定元素属于对角块并设置为(0,0)，否则设置为(0, None)，完全满足需求。
2. 冗余代码清理：代码中row_start、row_end、col_start、col_end四个变量未被使用，可直接删除，不影响功能。
3. 索引风格优化：Python习惯使用0起始索引，可将循环范围改为range(K*d)，省去索引减1的操作，逻辑更简洁：

   bnds = []
   for i in range(K*d):
       for j in range(K*d):
           block_row = i // K
           block_col = j // K
           bnds.append((0, 0) if block_row == block_col else (0, None))
   

4. 列表推导式简化：参考原代码的风格，可改成更高效的列表推导式：

   bnds = [(0, 0) if (i // K) == (j // K) else (0, None) for i in range(K*d) for j in range(K*d)]
   

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