Advanced constraint programming modeling

advancedModelingWeek2.py

from ortools.sat.python import cp_model

''' course: https://www.coursera.org/learn/advanced-modeling#about '''


def formTeam(model, booleanVarList, archerList, cavalryList, infantryList):

    varList = []
    for archer in archerList:
        varList.append(booleanVarList[archer])
    model.Add(sum(varList) >= 1)
    varList.clear()

    for cavalry in cavalryList:
        varList.append(booleanVarList[cavalry])
    model.Add(sum(varList) >= 2)
    varList.clear()

    for infantry in infantryList:
        varList.append(booleanVarList[infantry])
    model.Add(sum(varList) >= 2)

    model.Add(sum(booleanVarList) == 6)



def armySelectionProblem():

    model = cp_model.CpModel()

    archerList = [0, 1, 2, 3]
    cavalryList = [4, 5, 6, 7, 8, 9, 10]
    infantryList = [11, 12, 13, 14, 15, 16, 17]

    luiValueList = [2, 5, 6, 8, 9, 5, 8, 7, 7, 4, 6, 1, 2, 5, 6, 4, 8, 3]
    guanValueList = [9, 8, 4, 7, 6, 4, 5, 3, 5, 5, 7, 5, 8, 3, 9, 2, 4, 6]
    zhangValueList = [8, 4, 3, 3, 6, 2, 5, 5, 3, 2, 5, 4, 1, 1, 3, 4, 5, 6]

    numOfEliteSoldiers = len(archerList) + len(cavalryList) + len(infantryList)

    luiArmyVarList, guanArmyVarList, zhangArmyVarList = [],[],[]
    for soldier in range(numOfEliteSoldiers):
        luiArmyVarList.append(model.NewBoolVar("SelectedForLuiArmy|" + str(soldier)))
        guanArmyVarList.append(model.NewBoolVar("SelectedForGuanArmy|" + str(soldier)))
        zhangArmyVarList.append(model.NewBoolVar("SelectedForZhangArmy|" + str(soldier)))


    for soldier in range(numOfEliteSoldiers):
        model.Add(luiArmyVarList[soldier] + guanArmyVarList[soldier] + zhangArmyVarList[soldier] <= 1)

    formTeam(model, luiArmyVarList, archerList, cavalryList, infantryList)
    formTeam(model, guanArmyVarList, archerList, cavalryList, infantryList)
    formTeam(model, zhangArmyVarList, archerList, cavalryList, infantryList)

    objVarList = []
    for i, armyVar in enumerate(luiArmyVarList):
        objVarList.append(armyVar * luiValueList[i])
    for i, armyVar in enumerate(guanArmyVarList):
        objVarList.append(armyVar * guanValueList[i])
    for i, armyVar in enumerate(zhangArmyVarList):
        objVarList.append(armyVar * zhangValueList[i])

    model.Maximize(sum(objVarList))

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

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Maximum perceived value: %i' % solver.ObjectiveValue(), "\n")

        for armyVar in luiArmyVarList:
            if solver.Value(armyVar):
                print(armyVar.Name())
        for armyVar in guanArmyVarList:
            if solver.Value(armyVar):
                print(armyVar.Name())
        for armyVar in zhangArmyVarList:
            if solver.Value(armyVar):
                print(armyVar.Name())
    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#



def banditHouseCoveringProblem():
    model = cp_model.CpModel()

    size = 9
    numOfPoints = 3
    manhattanDist = 3
    costList = [[9, 9, 9, 9, 9, 9, 9, 9, 9],
                [8, 0, 8, 0, 8, 0, 8, 0, 9],
                [7, 7, 7, 7, 7, 7, 7, 8, 9],
                [6, 0, 6, 0, 6, 0, 7, 0, 9],
                [5, 5, 5, 5, 5, 6, 7, 8, 9],
                [4, 0, 4, 0, 5, 0, 7, 0, 9],
                [3, 3, 3, 4, 5, 6, 7, 8, 9],
                [2, 0, 3, 0, 5, 0, 7, 0, 9],
                [1, 2, 3, 4, 5, 6, 7, 8, 9]]
    costListFlattened = sum(costList, [])

    rowVarList = []
    colVarList = []
    indexVarList = []
    for point in range(numOfPoints):
        rowVarList.append(model.NewIntVar(1, size, "Row"+str(point)))
        colVarList.append(model.NewIntVar(1, size, "Col"+str(point)))

        indexVarList.append(model.NewIntVar(0, size*size, "indexFor"+str(point)))
        model.Add(indexVarList[point] == ((rowVarList[point] - 1) * size) + colVarList[point] - 1)

    model.AddAllDifferent(indexVarList)

    for i in range(numOfPoints):
        rowIsOddVar = model.NewBoolVar("RowIsOdd"+str(i))
        colIsOddVar = model.NewBoolVar("ColIsOdd" + str(i))
        model.AddBoolOr([rowIsOddVar, colIsOddVar])

        rowRemainderVar = model.NewBoolVar("RowRemainder"+str(i))
        model.AddModuloEquality(rowRemainderVar, rowVarList[i], 2)
        model.Add(rowRemainderVar == 1).OnlyEnforceIf(rowIsOddVar)

        colRemainderVar = model.NewBoolVar("ColRemainder"+str(i))
        model.AddModuloEquality(colRemainderVar, colVarList[i], 2)
        model.Add(colRemainderVar == 1).OnlyEnforceIf(colIsOddVar)


    hutList = [i*2 for i in range(1, size//2 + 1)]
    for x in hutList:
        for y in hutList:
            pointVar = model.NewIntVar(0, numOfPoints-1, "pointVarCoveringHutAt"+str(x)+str(y))
            tempRowVar = model.NewIntVar(1, size, "tempRow"+str(x)+str(y))
            tempColVar = model.NewIntVar(1, size, "tempCol"+str(x)+str(y))
            model.AddElement(pointVar, rowVarList, tempRowVar)
            model.AddElement(pointVar, colVarList, tempColVar)

            distanceRowVar = model.NewIntVar(-size, size, "distanceRow"+str(x)+str(y))
            distanceColVar = model.NewIntVar(-size, size, "distanceCol" + str(x) + str(y))
            model.Add(distanceRowVar == x - tempRowVar)
            model.Add(distanceColVar == y - tempColVar)

            absDistanceRowVar = model.NewIntVar(0, size, "absDistanceRow"+str(x)+str(y))
            absDistanceColVar = model.NewIntVar(0, size, "absDistanceCol" + str(x) + str(y))
            model.AddAbsEquality(absDistanceRowVar, distanceRowVar)
            model.AddAbsEquality(absDistanceColVar, distanceColVar)
            model.Add(absDistanceColVar + absDistanceRowVar <= manhattanDist)

    costVarList = []
    for point in range(numOfPoints):
        costVarList.append(model.NewIntVar(0, max(costListFlattened), "costFor"+str(point)))
        model.AddElement(indexVarList[point], costListFlattened, costVarList[point])
    model.Minimize(sum(costVarList))

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

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Cost: %i' % solver.ObjectiveValue(), "\n")

        rowVarSolList, colVarSolList = [], []
        for point in range(numOfPoints):
            rowVarSolList.append(solver.Value(rowVarList[point]))
            colVarSolList.append(solver.Value(colVarList[point]))

        print("Row: "+str(rowVarSolList))
        print("Col: "+str(colVarSolList))

    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#




def seatPlanningProblem():
    model = cp_model.CpModel()

    scholarList = ["S1", "S2", "S3", "S4", "S5", "S6", "S7", "S8", "S9", "S10", "S11",
                   "S12", "S13", "S14", "S15", "S16", "S17", "S18", "S19", "S20"]
    reputationList = [10, 7, 5, 6, 2, 10, 10, 10, 4, 3, 7, 8, 3, 9, 8, 6, 3, 5, 1, 7]
    maxReputation = max(reputationList)
    numOfScholars = len(scholarList)
    numOfTables = 5
    sizeOfTable = 5
    enemiesList = [[0,1]]
    friendsList = [[0,4]]
    keyGuestScholarList = [0, 5, 6, 7]

    seatVarList = []
    for scholar in range(numOfScholars):
        seatVarList.append(model.NewIntVar(0, numOfTables-1, "seatForScholar"+str(scholar)))

    for scholarPair in enemiesList:
        model.Add(seatVarList[scholarPair[0]] != seatVarList[scholarPair[1]])

    for scholarPair in friendsList:
        model.Add(seatVarList[scholarPair[0]] == seatVarList[scholarPair[1]])

    for i in range(len(keyGuestScholarList)-1):
        for j in range(i+1, len(keyGuestScholarList)):
            model.Add(seatVarList[i] != seatVarList[j])

    tableScholarVarList = []
    indexDict = {}
    counter = 0
    for scholar in range(numOfScholars):
        for table in range(numOfTables):
            tableScholarVarList.append(model.NewBoolVar("table"+str(table)+"scholar"+str(scholar)))
            indexDict[(table, scholar)] = counter
            counter += 1

    for scholar in range(numOfScholars):
        indexVar = model.NewIntVar(0, numOfTables*numOfScholars, "index"+str(scholar))
        model.Add(indexVar == numOfTables*scholar + seatVarList[scholar])
        model.AddElement(indexVar, tableScholarVarList, 1)

    scholarIndexVarList = []
    for scholar in range(numOfScholars):
        for table in range(numOfTables):
            scholarIndexVarList.append(model.NewIntVar(0, numOfScholars * 2, "scholarIndex"+str(table)+str(scholar)))
            model.Add(scholarIndexVarList[-1] == scholar + numOfScholars * (1 - tableScholarVarList[indexDict[(table, scholar)]]))

    minScholarIndexOnTableVarList = []
    for table in range(numOfTables):
        minScholarIndexOnTableVarList.append(model.NewIntVar(0, numOfScholars, "minScholarInTable" + str(table)))
        varList = []
        for scholar in range(numOfScholars):
            varList.append(scholarIndexVarList[indexDict[(table, scholar)]])
        model.AddMinEquality(minScholarIndexOnTableVarList[table], varList)

    # remove value symmetry
    for table in range(numOfTables-1):
        model.Add(minScholarIndexOnTableVarList[table] <= minScholarIndexOnTableVarList[table+1])


    tableUsedVarList = []
    for table in range(numOfTables):
        tableUsedVarList.append(model.NewBoolVar("tableUsed" + str(table)))
        varList = []
        for scholar in range(numOfScholars):
            varList.append(tableScholarVarList[indexDict[(table, scholar)]])
        model.Add(sum(varList) <= sizeOfTable * tableUsedVarList[table])
        model.Add(sum(varList) != 1)


    reputationDiffVarList = []
    for table in range(numOfTables):
        reputationDiffVarList.append(model.NewIntVar(0, maxReputation, "reputationDiff"+str(table)))
        maxReputationVar = model.NewIntVar(0, maxReputation, "maxReputation"+str(table))
        minReputationVar = model.NewIntVar(0, maxReputation, "minReputation"+str(table))
        varList = []
        for scholar in range(numOfScholars):
            varList.append(tableScholarVarList[indexDict[(table, scholar)]])
        model.AddMaxEquality(maxReputationVar, varList)
        model.AddMinEquality(minReputationVar, varList)
        model.Add(reputationDiffVarList[table] == maxReputationVar - minReputationVar)


    obj1Var = model.NewIntVar(0, numOfTables, "obj1Var")
    obj2Var = model.NewIntVar(0, numOfTables*maxReputation, "obj2Var")
    model.Add(obj1Var == sum(tableUsedVarList))
    model.Add(obj2Var == sum(reputationDiffVarList))
    model.Minimize(100*obj1Var + obj2Var)

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

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Obj1: '+str(solver.Value(obj1Var))+ ", Obj2: "+str(solver.Value(obj2Var))+"\n")

        for table in range(numOfTables):
            solList = []
            for scholar in range(numOfScholars):
                var = tableScholarVarList[indexDict[(table, scholar)]]
                if solver.Value(var):
                    solList.append(scholarList[scholar])
            print("Table "+str(table)+": "+str(solList))
    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#




def brigadeCreationProblem():
    model = cp_model.CpModel()

    numOfBattalion = 24
    archerBattalionIndices = [0, 1, 2, 3, 14]
    swordBattalionIndices = [2, 5, 7, 8, 14]
    shieldBattalionIndices = [0, 2, 14, 15, 17]

    sizeOfBrigade = 6
    numOfBrigade = numOfBattalion//sizeOfBrigade


    brigadeBattalionVarDict = {}
    for brigade in range(numOfBrigade):
        for battalion in range(numOfBattalion):
            brigadeBattalionVarDict[(brigade, battalion)] = model.NewBoolVar("brigade"+str(brigade)+"battalion"+str(battalion))


    for brigade in range(numOfBrigade):
        varList = []
        for battalion in range(numOfBattalion):
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) == sizeOfBrigade)


    for battalion in range(numOfBattalion):
        varList = []
        for brigade in range(numOfBrigade):
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) == 1)


    isBlackTurtleFormationVarList = []
    for brigade in range(numOfBrigade):
        varList = []
        isBlackTurtleFormationVarList.append(model.NewBoolVar("isBlackTurtleFormation" + str(brigade)))

        for battalion in archerBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 2).OnlyEnforceIf(isBlackTurtleFormationVarList[brigade])
        varList.clear()

        for battalion in shieldBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 2).OnlyEnforceIf(isBlackTurtleFormationVarList[brigade])


    isWhiteTigerFormationVarList = []
    for brigade in range(numOfBrigade):
        varList = []
        isWhiteTigerFormationVarList.append(model.NewBoolVar("isWhiteTigerFormation" + str(brigade)))

        for battalion in swordBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 3).OnlyEnforceIf(isWhiteTigerFormationVarList[brigade])


    isVemillionBirdFormationVarList = []
    for brigade in range(numOfBrigade):
        varList = []
        isVemillionBirdFormationVarList.append(model.NewBoolVar("isVermillionBirdFormation" + str(brigade)))

        for battalion in archerBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 1).OnlyEnforceIf(isVemillionBirdFormationVarList[brigade])
        varList.clear()

        for battalion in shieldBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 1).OnlyEnforceIf(isVemillionBirdFormationVarList[brigade])
        varList.clear()

        for battalion in swordBattalionIndices:
            varList.append(brigadeBattalionVarDict[(brigade, battalion)])
        model.Add(sum(varList) >= 2).OnlyEnforceIf(isVemillionBirdFormationVarList[brigade])


    eliteVarList = []
    for brigade in range(numOfBrigade):
        eliteVarList.append(model.NewBoolVar("isEliteBrigade"+str(brigade)))
        model.Add(eliteVarList[brigade] <= isBlackTurtleFormationVarList[brigade] +
                                           isWhiteTigerFormationVarList[brigade] +
                                           isVemillionBirdFormationVarList[brigade])

    model.Maximize(sum(eliteVarList))


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

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Maximum number of elite brigades: %i' % solver.ObjectiveValue(), "\n")

        for brigade, var in enumerate(eliteVarList):
            if solver.Value(var) > 0.0:
                print("Brigade: "+str(brigade)+ " is elite", end=", ")
                if solver.Value(isBlackTurtleFormationVarList[brigade]) > 0.0:
                    print("with black turtle formation", end=", ")
                if solver.Value(isWhiteTigerFormationVarList[brigade]) > 0.0:
                    print("with white tiger formation", end=", ")
                if solver.Value(isVemillionBirdFormationVarList[brigade]) > 0.0:
                    print("with vermillion bird formation")
                print()
    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#




if __name__ == '__main__':

    armySelectionProblem()
    banditHouseCoveringProblem()
    seatPlanningProblem()
    brigadeCreationProblem()