Is it possible to compute argmax with or-tools with just one integer variable?

Question

Part of a models objective is weighted by items in a scalar list.

I am solving for this by using a list of 0-1 range variables, and then using LinearExpr.ScalProd to weight the objective.

Is there a way to do this with just one integer variable ( other than the objective variable), where I can either use a lambda or some other mechanism to lookup the value of a variable in a table?

Here is some sample code that I have that works, though I am trying to determine it can be made more concise.

def argmax(
    model: cp_model.CpModel, values: List[int]
) -> Tuple[List[cp_model.IntVar], cp_model.IntVar]:
    objective_var = model.NewIntVar(0, 1000000, "objective_var")
    ret_vars = [model.NewIntVar(0, 1, "x(%i)" % i) for i in range(len(values))]
    model.Add(sum(ret_vars) == 1)
    model.Add(objective_var == cp_model.LinearExpr.ScalProd(ret_vars, values))

    return [ret_vars, objective_var]

Answer 1

This can be done using model.AddElement.

def AddElement(self, index, variables, target)

I found the documentation for AddElement to be a bit cryptic, so I'll try to curate it in what I'd think are simpler terms.

It states: Adds the element constraint: variables[index] == target.

• Index -> The variable you want to lookup with.
• Variables -> The table of lookup values.
• Target -> The value you want this to equal.

If anything, reading the C++ code makes it much simpler to understand.

Constraint CpModelBuilder::AddElement(IntVar index,
                                      absl::Span<const int64> values,
                                      IntVar target) {

Putting this all together, we get:

def argmax(
    model: cp_model.CpModel, values: List[int]
) -> Tuple[cp_model.IntVar, cp_model.IntVar]:
    objective_var = model.NewIntVar(0, 1000000, "objective_var")
    ret_var = model.NewIntVar(0, len(values) - 1, "x")
    model.AddElement(ret_var, values, objective_var)

    return [ret_var, objective_var]