How can I get row index in google or-tools? - or-tools

Question

I have a two-dimensional bool array, and I want to find the row index(index+1) of each column element that is not 0,If all the elements of the column are 0, the row index we get is 0. In addition, each column has at most 1 value that is not 0。How can I do it？
The following code cannot handle the case where one column is all zeros
model = cp_model.CpModel()
solver = cp_model.CpSolver()
shifts = {}
worker_shift={}
for i in range(6):
worker_shift[(i)]=model.NewIntVar(0, 6, "worker_shift(%i)" % (i))
for j in range(6):
shifts[(i,j)] = model.NewBoolVar("shifts(%i,%i)" % (i,j))
for j in range(6):
model.Add(sum(shifts[(i,j)] for i in range(6))<=1)
for j in range(6):
for i in range(6):
model.Add(worker_shift[(j)]==i+1).OnlyEnforceIf(shifts[(i,j)])
model.Add(shifts[(2,1)]==1)
r1= np.empty([6,6], dtype = int)
r2=np.empty([1,6], dtype = int)
status = solver.Solve(model)
print("status:",status)
for i in range(6):
r2[0,i]=solver.Value(worker_shift[(i)])
for j in range(6):
r1[i,j]=solver.Value(shifts[(i,j)])
print(r1)
print("----------------------")
print(r2)

Answer 1

Let's look at one column
[b1..bn]
if we create
pi == n - bi * (n - i)
Then using model.AddMinEquality(), min([pi]) computes the same think as your code.
Now let's introduce e a new Boolean variable.
we want
e <=> bool_or(pi)
We code it the usual way
for all i, pi => e (using model.AddImplication())
bool_or([b1, ..., bn, e.Not()]) (reverse implication, if all bi are false, then e is false.
Now, we can use the correct min() using a new IntVar pe[0..n]
model.Add(pe = e*n)
model.AddMinEquality(target, [p1, .., pn, pe])

Answer 2

thanks to #Laurent Perron,#Stradivari ！
I studied the introduction of or-tools in #Stradivari's blog(https://www.xiang.dev/cp-sat/#iff-equivalence-boolean-product), I use following code:
for j in range(6):
model.Add(worker_shift[(j)]==0).OnlyEnforceIf([shifts[(i,j)].Not() for i in range(6) ])
for i in range(6):
model.Add(worker_shift[(j)]==i+1).OnlyEnforceIf(shifts[(i,j)])
It seems to work well !
The complete code is as follows：
model = cp_model.CpModel()
solver = cp_model.CpSolver()
shifts = {}
worker_shift={}
for i in range(6):
worker_shift[(i)]=model.NewIntVar(0, 6, "worker_shift(%i)" % (i))
for j in range(6):
shifts[(i,j)] = model.NewBoolVar("shifts(%i,%i)" % (i,j))
for j in range(6):
model.Add(sum(shifts[(i,j)] for i in range(6))<=1)
for j in range(6):
model.Add(worker_shift[(j)]==0).OnlyEnforceIf([shifts[(i,j)].Not() for i in range(6) ])
for i in range(6):
model.Add(worker_shift[(j)]==i+1).OnlyEnforceIf(shifts[(i,j)])
model.Add(worker_shift[(2)]==0)
model.Add(worker_shift[(0)]==0)
model.Add(worker_shift[(3)]==0)
model.Add(shifts[(2,1)]==1)
model.Add(shifts[(3,4)]==1)
r1= np.empty([6,6], dtype = int)
r2=np.empty([1,6], dtype = int)
status = solver.Solve(model)
print("status:",status)
for i in range(6):
r2[0,i]=solver.Value(worker_shift[(i)])
for j in range(6):
r1[i,j]=solver.Value(shifts[(i,j)])
print(r1)
print("----------------------")
print(r2)