how to make K different elements in bool array with or-tools?

467 views Asked by At

I want make K different elements in bool array,I use code: model.Add(len(set([shifts[(i)] for i in range(10)]))==4) ,but it not work! How can I do this?

model = cp_model.CpModel()
solver = cp_model.CpSolver()

shifts = {}
ones = [model.NewBoolVar("") for _ in range(10)]
for i in range(10):
    shifts[(i)] = model.NewIntVar(0, 10, "shifts(%i)" % i)
    
for i in range(10):
    model.Add(shifts[(i)] >0).OnlyEnforceIf(ones[(i)])
    model.Add(shifts[(i)] == 0).OnlyEnforceIf(ones[(i)].Not())
    
model.Add(sum(ones[(i)] for i in range(10)) == 5)

# I want make 4 different but it not work!
#model.Add(len(set([shifts[(i)] for i in range(10)]))==4)

status = solver.Solve(model)
print("status:",status)

res=[]
for i in range(10):
        res.append(solver.Value(shifts[(i)]))
print(res)
1

There are 1 answers

1
Stradivari On

Encode your integers using booleans and add another boolean for each value to mark it as used

from ortools.sat.python import cp_model

model = cp_model.CpModel()
solver = cp_model.CpSolver()

shifts = {}
used = [model.NewBoolVar("") for j in range(10)]
for i in range(10):
    for j in range(10):
        shifts[i, j] = model.NewBoolVar(f"shifts({i}, {j})")
        # shifts[i,j] => used[j]
        model.AddImplication(shifts[i, j], used[j])
    model.Add(sum(shifts[i, j] for j in range(10)) == 1)

model.Add(sum(shifts[i, 0].Not() for i in range(10)) == 5)

for j in range(10):
    # all(shifts[_, j] == 0) => used[j].Not()
    model.AddBoolOr([shifts[i, j] for i in range(10)] + [used[j].Not()])

model.Add(sum(used) == 4)

status = solver.Solve(model)
print("status:", status)

res = []
for i in range(10):
    for j in range(10):
        if solver.Value(shifts[i, j]):
            res.append(j)
print(res)