如何用Python OR-Tools约束二维数组指定位置为特定数值

我有一个二维列表list_of_numbers，每个子列表包含1个OR-Tools模型的IntVar整数变量，取值范围为0到10。另有一个二维列表list_of_positions，需要基于它编写约束，使list_of_numbers中的元素满足以下规则：

• list_of_positions的索引对应目标数字（比如索引0对应数字0，索引1对应数字1）
• 每个子列表中的数值是该数字必须出现的list_of_numbers的位置索引
• 空列表表示对应数字不能出现在list_of_numbers的任何位置

举例：当list_of_positions为[[1],[2],[0,3],[]]时：

• 数字0必须在list_of_numbers的索引1位置
• 数字1必须在list_of_numbers的索引2位置
• 数字2必须在list_of_numbers的索引0和3位置
• 数字3不能出现在任何位置
  最终list_of_numbers应变为[[2],[0],[1],[2]]

我写的代码中约束部分有误，请求修正：

from ortools.sat.python import cp_model

model = cp_model.CpModel()

#prints the solution at the end
def print_solution(list_of_numbers, solver):
    for i in range(0, len(list_of_numbers)):
        for j in range(0, 1):
            print(f'{solver.Value(list_of_numbers[i][j]):2} ', end = '')

        print()

#creates a list of 4 int variables that can be from 0 to 10
list_of_numbers = []
for i in range(0, 4):
    number = []
    for j in range(0, 1):
        number.append(model.NewIntVar(0, 10, f'list_of_numbers{i}_{j}'))
    list_of_numbers.append(number)

list_of_positions =  [[1],[2],[0,3],[]]

#this is where i tried writing a constraint but im terrible with or-tools and i know this is way off
####################################
for i in range(0,len(list_of_positions)):
    for j in range(0,len(list_of_positions[i])):
        model.Add(list_of_numbers[i] == list_of_positions[i][j])
#####################################

solver = cp_model.CpSolver()
status = solver.Solve(model)

if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
    print_solution(list_of_numbers, solver)
else:
    print('unsat')

原代码的约束逻辑完全搞反了对应关系：你把list_of_positions的索引（目标数字）当成了list_of_numbers的位置索引，错误地给list_of_numbers[i]添加等于list_of_positions[i][j]的约束，这和需求完全不符。

正确的逻辑应该是：

• list_of_positions[d]中的d是目标数字，子列表里的p是list_of_numbers的位置索引
• 若list_of_positions[d]为空，所有位置都不能等于d
• 若list_of_positions[d]非空，只有指定位置p可以等于d，其他位置都不能等于d

修正后的代码如下：

from ortools.sat.python import cp_model

model = cp_model.CpModel()

# 打印解
def print_solution(list_of_numbers, solver):
    for i in range(len(list_of_numbers)):
        print(f'{solver.Value(list_of_numbers[i][0]):2}')

# 创建4个取值0-10的整数变量（二维列表，每个子列表含1个变量）
list_of_numbers = []
for i in range(4):
    var = model.NewIntVar(0, 10, f'pos_{i}')
    list_of_numbers.append([var])

list_of_positions = [[1], [2], [0, 3], []]

# 修正后的约束逻辑
####################################
# 遍历每个目标数字d（对应list_of_positions的索引）
for d in range(len(list_of_positions)):
    target_positions = list_of_positions[d]
    if not target_positions:
        # 空列表：所有位置都不能等于当前数字d
        for p in range(len(list_of_numbers)):
            model.Add(list_of_numbers[p][0] != d)
    else:
        # 非空列表：指定位置必须等于d，其余位置不能等于d
        for p in range(len(list_of_numbers)):
            if p in target_positions:
                model.Add(list_of_numbers[p][0] == d)
            else:
                model.Add(list_of_numbers[p][0] != d)
#####################################

solver = cp_model.CpSolver()
status = solver.Solve(model)

if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
    print_solution(list_of_numbers, solver)
else:
    print('unsat')

运行后会输出符合预期的结果：

2
 0
 1
 2

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