ORtools assigning either one regular or three special classes to each student

Question

In my problem, I have students who need to be assigned classes. Either they get assigned one regular class, or three special classes. I have lists students and classes, corresponding to the student and class schedules:

students = [s0,s1,s2,s3,s4,s5,s6,s7,s8,s9]
classes = [sp0,sp1,sp2,sp3,sp4,sp5,sp6,r0,r1,r2,r3,r4,r5,r6]

where the first 7 classes are special classes (a student is assigned three of these), and the last 7 are regular classes (a student is assigned one of these). The basic setup looks like this:

def main():
    # Data.
    num_students = len(students)
    num_classes = len(classes)
    all_students = range(num_students)
    all_classes = range(num_classes)
    k = 7
    special_classes = range(k) # k = number of special classes
    regular_classes = range(k,num_classes)

    # Creates the model.
    model = cp_model.CpModel()

    # Creates scheduling variables.
    # sched[(c, s)]: student 'c' is assigned to class 's'
    # 1 if true, 0 if not
    sched = {}
    for s in all_students:
        for c in all_classes:
                sched[(c,s)] = model.NewBoolVar('shift_c%is%i' % (c, s))

But now I need a constraint so that each student is either assigned one regular class or three special classes. I was thinking something like:

for s in all_students:

        # either a student is assigned 3 special classes or 1 regular
        model.AddBoolOr([sum(sched[(c,s)] for c in other_classes)==1,
                        sum(sched[(c,s)] for c in special_classes)==3])

        # either a student is assigned 0 special classes or 3
        model.AddBoolOr([sum(sched[(c,s)] for c in special_classes)==0,
                        sum(sched[(c,s)] for c in special_classes)==3])

        # either a student is assigned 1 regular class or 0
        model.AddBoolOr([sum(sched[(c,s)] for c in other_classes)==1,
                        sum(sched[(c,s)] for c in other_classes)==0])

but I realize I can't use AddBoolOr like this.

How do you think I could solve this problem?

Any help appreciated!

Answer 1

You can use an intermediate boolean for each student:

for s in all_students:
    regular = model.NewBoolVar('reg_%i' % s)
    model.Add(sum(sched[(c,s)] for c in regular_classes) == 1).OnlyEnforceIf(regular)
    model.Add(sum(sched[(c,s)] for c in special_classes) == 0).OnlyEnforceIf(regular)

    model.Add(sum(sched[(c,s)] for c in regular_classes) == 0).OnlyEnforceIf(regular.Not())
    model.Add(sum(sched[(c,s)] for c in special_classes) == 3).OnlyEnforceIf(regular.Not())

Please take a look at https://github.com/google/or-tools/blob/stable/ortools/sat/doc/channeling.md to learn more about modeling in ortools.