I will be grateful to anyone who can help me write some Python code to enumerate the 21×2×3 arrays, indexed with i, j and k, which are two-thirds filled with 0's and one-third filled with the values 'Ava', 'Bob', 'Joe', 'Mia', 'Sam', 'Tom', 'Zoe' in such a way that:
fixed the index i you have exactly two empty 2-tuples and one 2-tuple with different non-zero values;
fixed the index k you have exactly fourteen empty 2-tuples and seven 2-tuple with different non-zero values;
fixed the indexes j and k you have a 21-tuple with fourteen zero values and exactly one occurrence of each of the non-zero values, respecting the following constraints:
a) "Ava" can appear only in a row with index 0, 1, 4, 6, 10, 11, 13, 14, 15, 19 or 20;
b) "Bob" can appear only in a row with index 2, 3, 5, 7, 8, 9, 12, 16, 17 or 18;
c) "Joe" can appear only in a row with index 2, 4, 5, 7, 8, 10, 14, 15, 18 or 20;
d) "Mia" can appear only in a row with index 0, 1, 3, 6, 9, 12, 13, 16, 17 or 19;
e) "Sam" can appear only in a row with index 1, 2, 7, 9, 15, 17 or 20;
f) "Tom" can appear only in a row with index 0, 3, 8, 10, 12, 16 or 19;
g) "Zoe" can appear only in a row with index 4, 5, 6, 11, 13, 14 or 18.
As a result I would like to obtain something like this:
[ 0 0 [Tom Mia [ 0 0
0 0 Ava Sam 0 0
0 0 Sam Bob 0 0
0 0 Bob Tom 0 0
0 0 0 0 Joe Zoe
0 0 Joe Zoe 0 0
0 0 0 0 Zoe Ava
Joe Sam 0 0 0 0
0 0 0 0 Tom Bob
0 0 0 0 Mia Sam
Tom Ava 0 0 0 0
Ava Zoe 0 0 0 0
Bob Mia 0 0 0 0
0 0 Mia Ava 0 0
0 0 Zoe Joe 0 0
Sam Joe 0 0 0 0
0 0 0 0 Bob Tom
0 0 0 0 Sam Mia
Zoe Bob 0 0 0 0
Mia Tom 0 0 0 0
0 0 ] 0 0 ] Ava Joe]
Rows represent school classes, columns represent school terms (there are 2 of them), tubes represent class days (there are 3 of them: Monday, Wednesday and Friday). So the first horizontal slice of the above solution means that class 1A has lesson only on Wednesday, in the the first term with teacher Tom and in the second term with teacher Mia. (Teachers can only work in some classes and not in others.)
Thanks in advance!
Update n. 1
As a starting point, I tried to attack the following toy problem:
Enumerate all arrays with a given number of rows and 3 columns which are two-thirds filled with "0" and one-third filled with "1" in such a way that summing the values in each row you always get 1 and summing the values in each column you always get rows / 3.
Finally, after struggling a bit, I think I managed to get a solution with the following code, that I kindly ask you to correct or improve. (I have set rows = 6 because the number of permutations of the obvious solution is 6!/(2!*2!*2!) = 90, whereas setting rows = 21 I would have got 21!/(7!*7!*7!) = 399,072,960 solutions.)
from ortools.sat.python import cp_model
# Create model.
model = cp_model.CpModel()
# Create variables.
rows = 6
columns = 3
x = []
for i in range(rows):
x.append([model.NewBoolVar(f'x[{i}][{j}]') for j in range(columns)])
# Add constraints.
for i in range(rows):
model.Add(sum(x[i]) == 1)
# Uncomment the following four lines of code if you want to solve the slightly more general problem that asks to enumerate
# all boolean arrays, with a given number of rows and columns, filled in such a way that summing the values in each
# row you always get 1 and summing the values in each column you always get no more than the ceiling of (rows / columns).
# if rows % columns != 0:
# for j in range(columns):
# model.Add(sum(x[i][j] for i in range(rows)) <= rows // columns + 1)
# else:
for j in range(columns):
model.Add(sum(x[i][j] for i in range(rows)) == rows // columns)
class MyPrintedSolution():
def __init__(self, sol, sol_number):
self.sol = sol
self.sol_number = sol_number
def PrintReadableTable(self):
print(f'Solution {self.sol_number}, printed in readable form:')
counter = 0
for v in self.sol:
if counter % columns != columns-1:
print(v, end = ' ')
else:
print(v)
counter += 1
print()
def PrintRawSolution(self):
print(f'Solution {self.sol_number}, printed in raw form:')
counter = 0
for v in self.sol:
print(f'{v}', end = '')
counter += 1
print('\n')
class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
def __init__(self, variables, limit):
cp_model.CpSolverSolutionCallback.__init__(self)
self.__variables = variables
self.__solution_count = 0
self.__solution_limit = limit
def solution_count(self):
return self.__solution_count
def on_solution_callback(self):
self.__solution_count += 1
solution = [self.Value(v) for v in self.__variables]
myprint = MyPrintedSolution(solution, self.__solution_count)
myprint.PrintReadableTable()
# myprint.PrintRawSolution()
if self.__solution_count >= self.__solution_limit:
print(f'Stop search after {self.__solution_limit} solutions')
self.StopSearch()
# Create solver and solve model.
solver = cp_model.CpSolver()
# solver.parameters.num_workers = 16 # Solver works better with more workers. (At least 8, 16 if enough cores.)
# solver.parameters.log_search_progress = True
solver.parameters.enumerate_all_solutions = True
# solver.parameters.max_time_in_seconds = 10.0
solution_limit = 100000
solution_printer = VarArraySolutionPrinter([x[i][j] for i in range(rows) for j in range(columns)], solution_limit)
solver.Solve(model, solution_printer)
Update n. 2
Following #Christopher Hamkins' initial roadmap and subsequent precious suggestions, I think I finally got what I wanted, using the following code (although I am of course always open to corrections or further suggestions).
from ortools.sat.python import cp_model
# Create model.
model = cp_model.CpModel()
# Create variables.
classes = 21 # indexed with "i", but one could as well have chosen "c"
terms = 2 # indexed with "j", but one could as well have chosen "t"
days = 3 # indexed with "k", but one could as well have chosen "d"
persons = 8 # indexed with "p"
persons_names = [' 0 ', 'Ava', 'Bob', 'Joe', 'Mia', 'Sam', 'Tom', 'Zoe']
classes_names = ['1A', '1B', '1C', '1D', '1E', '1F', '1G', '2A', '2B', '2C', '2D', '2E', '2F', '2G', '3A', '3B', '3C', '3D', '3E', '3F', '3G']
classes_p = [[] for _ in range(persons)]
classes_p[0] = list(range(classes))
classes_p[1] = [0, 1, 4, 6, 10, 11, 13, 14, 15, 19, 20] # list of classes in which person 1 can work
classes_p[2] = [2, 3, 5, 7, 8, 9, 12, 16, 17, 18] # list of classes in which person 2 can work
classes_p[3] = [2, 4, 5, 7, 8, 10, 14, 15, 18, 20] # list of classes in which person 3 can work
classes_p[4] = [0, 1, 3, 6, 9, 12, 13, 16, 17, 19] # list of classes in which person 4 can work
classes_p[5] = [1, 2, 7, 9, 15, 17, 20] # list of classes in which person 5 can work
classes_p[6] = [0, 3, 8, 10, 12, 16, 19] # list of classes in which person 6 can work
classes_p[7] = [4, 5, 6, 11, 13, 14, 18] # list of classes in which person 7 can work
x = {}
for i in range(classes):
for j in range(terms):
for k in range(days):
for p in range(persons):
x[i, j, k, p] = model.NewBoolVar(f'x[{i}, {j}, {k}, {p}]')
# Add constraints.
"""
For all i, j, k constrain the sum of x[i, j, k, p] over p in the range of people to be equal to 1,
so exactly nobody or one person is selected at a given slot.
"""
for i in range(classes):
for j in range(terms):
for k in range(days):
model.Add(sum(x[i, j, k, p] for p in range(persons)) == 1)
"""
For all i constrain the sum of x[i, j, k, p] over all j, k, p in their respective ranges (except p = 0)
to be exactly equal to 2, so exactly two people are in a given row.
"""
for i in range(classes):
model.Add(sum(x[i, j, k, p] for j in range(terms) for k in range(days) for p in range(1, persons)) == 2)
"""
For all i, k, and for p = 0, add the implications
x[i, 0, k, 0] == x[i, 1, k, 0]
"""
for i in range(classes):
for k in range(days):
model.Add(x[i, 0, k, 0] == x[i, 1, k, 0])
"""
For all i, p (except p = 0), constrain the sum of x[i, j, k, p] over all j and k
to be at most 1.
"""
for i in range(classes):
for p in range(1, persons):
model.Add(sum(x[i, j, k, p] for j in range(terms) for k in range(days)) <= 1)
# for k in range(days): # Equivalent alternative to the previous line of code
# model.AddBoolOr([x[i, 0, k, p].Not(), x[i, 1, k, p].Not])
"""
For all j, k constrain the sum of x[i, j, k, p] over all i, p in their respective ranges (except p = 0)
to be exactly equal to 7, so exactly seven people are in a given column.
"""
for j in range(terms):
for k in range(days):
model.Add(sum(x[i, j, k, p] for i in range(classes) for p in range(1, persons)) == 7)
"""
For all j, k, p (except p = 0) constrain the sum of x[i, j, k, p] over all i
to be exactly equal to 1, so each person appears exactly once in the column.
"""
for j in range(terms):
for k in range(days):
for p in range(1, persons):
model.Add(sum(x[i, j, k, p] for i in range(classes)) == 1)
"""
For all j and k, constrain x[i, j, k, p] == 0 for the row i in which each person p can't appear.
"""
for p in range(persons):
for i in enumerate(set(range(classes)) - set(classes_p[p])):
for j in range(terms):
for k in range(days):
model.Add(x[i[1], j, k, p] == 0)
class MyPrintedSolution():
def __init__(self, sol, sol_number):
self.sol = sol
self.sol_number = sol_number
def PrintReadableTable1(self):
print(f'Solution {self.sol_number}, printed in first readable form:')
print(' | Mon | Wed | Fri ')
print(' Cl | Term1 Term2 | Term1 Term2 | Term1 Term2')
print('----------------------------------------------------', end='')
q = [_ for _ in range(8)] + [_ for _ in range(24, 32)] + [_ for _ in range(8, 16)] + [_ for _ in range(32, 40)] + [_ for _ in range(16, 24)] + [_ for _ in range(40, 48)]
r = []
for i in range(21):
r += [n+48*i for n in q]
shuffled_sol = [self.sol[m] for m in tuple(r)]
counter = 0
for w in shuffled_sol:
if (counter % (persons * days * terms)) == 0:
print('\n ', classes_names[counter // (terms * days * persons)], sep='', end=' |')
if w:
print(' ', persons_names[counter % persons], sep='', end=' ')
counter += 1
print('\n')
def PrintReadableTable2(self):
print(f'Solution {self.sol_number}, printed in second readable form:')
print(' Cl | Term1 Term2 ')
print(' Cl | Mon Wed Fri Mon Wed Fri ')
print('----------------------------------------', end = '')
counter = 0
for v in self.sol:
if (counter % (persons * days * terms)) == 0:
print('\n ', classes_names[counter // (terms * days * persons)], sep = '', end = ' |')
if v:
print(' ', persons_names[counter % persons], sep = '', end = ' ')
counter += 1
print('\n')
def PrintRawSolution(self):
print(f'Solution {self.sol_number}, printed in raw form:')
counter = 0
for v in self.sol:
print(f'{v}', end = '')
counter += 1
print('\n')
class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
def __init__(self, variables, limit):
cp_model.CpSolverSolutionCallback.__init__(self)
self.__variables = variables
self.__solution_count = 0
self.__solution_limit = limit
def solution_count(self):
return self.__solution_count
def on_solution_callback(self):
self.__solution_count += 1
solution = [self.Value(v) for v in self.__variables]
myprint = MyPrintedSolution(solution, self.__solution_count)
myprint.PrintReadableTable1()
# myprint.PrintReadableTable2()
# myprint.PrintRawSolution()
if self.__solution_count >= self.__solution_limit:
print(f'Stop search after {self.__solution_limit} solutions')
self.StopSearch()
# Create solver and solve model.
solver = cp_model.CpSolver()
# solver.parameters.num_workers = 16 # Solver works better with more workers. (At least 8, 16 if enough cores.)
# solver.parameters.log_search_progress = True
solver.parameters.enumerate_all_solutions = True
# solver.parameters.max_time_in_seconds = 10.0
solution_limit = 20
solution_printer = VarArraySolutionPrinter([x[i, j, k, p] for i in range(classes) for j in range(terms) for k in range(days) for p in range(persons)], solution_limit)
status = solver.Solve(model, solution_printer)
Update n. 3
#AirSquid proposed a solution using PuLP which is to me almost as valuable as the one using CP-SAT. It provides only one solution at a time, but (it has other advantages and) one can always get around this by adding some ad hoc further constraints, for example to see a different solution with a certain person in a specific position.
Your "toy" problem is definitely going in the right direction.
For your actual problem, try making a 21×2×3x8 array x indexed with i, j, k and p (for person) of BoolVar's. The last index represents the person, it will need 0 to represent "nobody" and for the rest Ava = 1, Bob = 2, etc., so its max value will be one more than the number of people. If the variable X[i,j,k,p] is true (1) it means that the given person p is present at the index i, j, k. If X[i,j,k,0] is true, it means a 0 = nobody is present at i, j, k.
For all i, j, k, constrain the sum of x[i, j, k, p] for p in the range of people to be equal to 1, so exactly nobody or one person is selected at a given slot.
For point 1: fixed the index i you have exactly two empty 2-tuples and one 2-tuple with different non-zero values:
For all i constrain the sum of x[i, j, k, p] for all j, k, p in their respective ranges (except p = 0) to be exactly equal to 2, so exactly two people are in a given row.
For all i, k, and for p = 0, add the implications
x[i, 0, k, 0] == x[i, 1, k, 0]
This will ensure that if one of the pair is 0, so is the other.
For all i, k and p except p = 0, add the implications
x[i, 0, k, p] implies x[i, 1, k, p].Not and
x[i, 1, k, p] implies x[i, 0, k, p].Not
(Actually one of these alone should be sufficient)
You can directly add an implication with the AddImplication(self, a, b) method, or you can realize that "a implies b" means the same thing as "b or not a" and add the implication with the AddBoolOr method. For the first implication, with x[i, 0, k, p] as a, and x[i, 1, k, p].Not as b, therefore adding:
AddBoolOr([x[i, 0, k, p].Not(), x[i, 1, k, p].Not])
Note that both variables are negated with Not in the expression.
Since the other implication assigns x[i, 1, k, p] as a, and x[i, 0, k, p].Not as b, the resulting expression is exactly the same
AddBoolOr([x[i, 0, k, p].Not(), x[i, 1, k, p].Not])
so you only need to add it once.
This will ensure a tuple will consist of two different people.
Alternative formulation of the last part:
For all i and p except p = 0, constraint the sum of x[i, j, k, p] for all j and k to be exactly equal to 1.
For point 2: fixed the index k you have exactly fourteen empty 2-tuples and seven 2-tuple with different non-zero values;
For all j and k constrain the sum of x[i, j, k, p] for all i and p (except p=0) in their respective ranges to be exactly equal to 7, so exactly seven people are in a given column.
For all j, k, and p (except p = 0) constrain the sum of x[i, j, k, p] over all i to be exactly equal to 1, so each person appears exactly once in the column (that is, once for each value of the indices j and k, for some value of i).
For point 3:
For all j and k, Constrain x[i, j, k, p] == 0 for the row i in which each person p can't appear.
Let us know how it works.
You're taking a pretty big swing if you are new to the trifecta of python, linear programming, and pulp, but the problem you describe is very doable...perhaps the below will get you started. It is a smaller example that should work just fine for the data you have, I just didn't type it all in.
A couple notes:
The below is a linear program. It is "naturally integer" as coded, preventing the need to restrict the domain of the variables to integers, so it is much easier to solve. (A topic for you to research, perhaps).
You could certainly code this up as a constraint problem as well, I'm just not as familiar. You could also code this up like a matrix as you are doing with i, j, k, but most frameworks allow more readable names for the sets.
The teaching day M/W/F is arbitrary and not linked to anything else in the problem, so you can (externally to the problem), just pick 1/3 of the assignments per day from the solution for each course & term.
The transition from the verbiage to the constraint formulation is most of the magic in linear programming and you'd be well suited with an introductory text if you continue along!
Code:
# teacher assignment
import pulp
from itertools import chain
# some data...
teach_days = {'M', 'W', 'F'}
terms = {'Spring', 'Fall'}
courses = {'Math 101', 'English 203', 'Physics 201'}
legal_asmts = { 'Bob': {'Math 101', 'Physics 201'},
'Ann': {'Math 101', 'English 203'},
'Tim': {'English 203'},
'Joe': {'Physics 201'}}
# quick sanity check
assert courses == set.union(*chain(legal_asmts.values())), 'course mismatch'
# set up the problem
prob = pulp.LpProblem('teacher_assignment', pulp.LpMaximize)
# make a 3-tuple index of the term, class, teacher
idx = [(term, course, teacher) for term in terms for course in courses for teacher in legal_asmts.keys()]
assn = pulp.LpVariable.dicts('assign', idx, cat=pulp.LpContinuous, lowBound=0)
# OBJECTIVE: teach as many courses as possible within constraints...
prob += pulp.lpSum(assn)
# CONSTRAINTS
# teach each class no more than once per term
for term in terms:
for course in courses:
prob += pulp.lpSum(assn[term, course, teacher] for teacher in legal_asmts.keys()) <= 1
# each teacher no more than 1 course per term
for term in terms:
for teacher in legal_asmts.keys():
prob += pulp.lpSum(assn[term, course, teacher] for course in courses) <= 1
# each teacher can only teach within legal assmts, and if legal, only teach it once
for teacher in legal_asmts.keys():
for course in courses:
if course in legal_asmts.get(teacher):
prob += pulp.lpSum(assn[term, course, teacher] for term in terms) <= 1
else: # it is not legal assignment
prob += pulp.lpSum(assn[term, course, teacher] for term in terms) <= 0
prob.solve()
#print(prob)
# Inspect results...
for i in idx:
if assn[i].varValue: # will be true if value is non-zero
print(i, assn[i].varValue)
Output:
Coin0008I MODEL read with 0 errors
Option for timeMode changed from cpu to elapsed
Presolve 16 (-10) rows, 12 (-12) columns and 32 (-40) elements
Perturbing problem by 0.001% of 1 - largest nonzero change 0.00010234913 ( 0.010234913%) - largest zero change 0
0 Obj -0 Dual inf 11.99913 (12)
10 Obj 5.9995988
Optimal - objective value 6
After Postsolve, objective 6, infeasibilities - dual 0 (0), primal 0 (0)
Optimal objective 6 - 10 iterations time 0.002, Presolve 0.00
Option for printingOptions changed from normal to all
Total time (CPU seconds): 0.00 (Wallclock seconds): 0.00
('Spring', 'Math 101', 'Bob') 1.0
('Spring', 'Physics 201', 'Joe') 1.0
('Spring', 'English 203', 'Ann') 1.0
('Fall', 'Math 101', 'Ann') 1.0
('Fall', 'Physics 201', 'Bob') 1.0
('Fall', 'English 203', 'Tim') 1.0
EDIT: corrected form of problem
Misunderstood part of the problem statement. The below is fixed. Needed to introduce a binary indicator variable for the day assignment per form and had some fun with tabulate.
Using an LP has the advantage that (with the included obj statement) it will do the best possible within the constraints to teach as much as possible, even if there is a teacher shortage, where a CP will not. A CP on the other hand can enumerate all the combos that satisfy the constraints, the LP cannot.
Code
# teacher assignment
import pulp
from tabulate import tabulate
# some data...
teach_days = {'M', 'W', 'F'}
terms = {'Spring', 'Fall'}
forms = list(range(20))
teach_capable = { "Ava" : [ 0, 1, 4, 6, 10, 11, 13, 14, 15, 19, 20],
"Bob" : [ 2, 3, 5, 7, 8, 9, 12, 16, 17, 18],
"Joe" : [ 2, 4, 5, 7, 8, 10, 14, 15, 18, 20],
"Mia" : [ 0, 1, 3, 6, 9, 12, 13, 16, 17, 19],
"Sam" : [ 1, 2, 7, 9, 15, 17, 20],
"Tom" : [ 0, 3, 8, 10, 12, 16, 19],
"Zoe" : [ 4, 5, 6, 11, 13, 14, 18],}
# set up the problem
prob = pulp.LpProblem('teacher_assignment', pulp.LpMaximize)
# make a 4-tuple index of the day, term, class, teacher
idx = [(day, term, form, teacher)
for day in teach_days
for term in terms
for form in forms
for teacher in teach_capable.keys()]
# variables
assn = pulp.LpVariable.dicts('assign', idx, cat=pulp.LpContinuous, lowBound=0)
form_day = pulp.LpVariable.dicts('form-day',
[(form, day) for form in forms for day in teach_days],
cat=pulp.LpBinary) # inidicator for which day the form uses
# OBJECTIVE: teach as many courses as possible within constraints...
prob += pulp.lpSum(assn)
# CONSTRAINTS
# 1. Teach each form on no more than 1 day
for form in forms:
prob += pulp.lpSum(form_day[form, day] for day in teach_days) <= 1 # limit to 1 day per form
for form in forms:
for day in teach_days:
for term in terms:
# no more than 1 assignment, if this day is the designated "form-day"
prob += pulp.lpSum(assn[day, term, form, teacher] for teacher in teach_capable.keys()) \
<= form_day[form, day]
# 2. Each teacher can only teach within legal assmts, and limit them to teaching that form once
for teacher in teach_capable.keys():
for form in forms:
if form in teach_capable.get(teacher):
prob += pulp.lpSum(assn[day, term, form, teacher] for day in teach_days for term in terms) <= 1
else: # it is not legal assignment
prob += pulp.lpSum(assn[day, term, form, teacher] for day in teach_days for term in terms) <= 0
# 3. Each teacher can only teach on once per day per term
for teacher in teach_capable.keys():
for term in terms:
for day in teach_days:
prob += pulp.lpSum(assn[day, term, form, teacher] for form in forms) <= 1
prob.solve()
print("Status = %s" % pulp.LpStatus[prob.status])
#print(prob)
# gather results...
selections = []
for i in idx:
if assn[i].varValue: # will be true if value is non-zero
selections.append(i)
#print(i, assn[i].varValue)
# Let's try to make some rows for tabulate... hacky but fun
def row_index(label):
"""return the form, column index, and name"""
col = 1
if 'W' in label: col += 2
elif 'F' in label: col += 4
if 'Fall' in label: col += 1
return label[2], col, label[-1]
headers = ['Form', 'Mon-1', 'Mon-2', 'Wed-1', 'Wed-2', 'Fri-1', 'Fri-2']
row_data = [[f,'','','','','',''] for f in forms]
for selection in selections:
form, col, name = row_index(selection)
row_data[form][col] = name
print(tabulate(row_data, headers=headers, tablefmt='grid'))
Output:
Status = Optimal
+--------+---------+---------+---------+---------+---------+---------+
| Form | Mon-1 | Mon-2 | Wed-1 | Wed-2 | Fri-1 | Fri-2 |
+========+=========+=========+=========+=========+=========+=========+
| 0 | | | Ava | Tom | | |
+--------+---------+---------+---------+---------+---------+---------+
| 1 | Mia | Sam | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 2 | | | | | Sam | Joe |
+--------+---------+---------+---------+---------+---------+---------+
| 3 | | | | | Bob | Mia |
+--------+---------+---------+---------+---------+---------+---------+
| 4 | Ava | Zoe | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 5 | Zoe | Joe | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 6 | | | Mia | Zoe | | |
+--------+---------+---------+---------+---------+---------+---------+
| 7 | | | Joe | Sam | | |
+--------+---------+---------+---------+---------+---------+---------+
| 8 | Bob | Tom | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 9 | | | Sam | Mia | | |
+--------+---------+---------+---------+---------+---------+---------+
| 10 | | | | | Ava | Tom |
+--------+---------+---------+---------+---------+---------+---------+
| 11 | | | Zoe | Ava | | |
+--------+---------+---------+---------+---------+---------+---------+
| 12 | | | Tom | Bob | | |
+--------+---------+---------+---------+---------+---------+---------+
| 13 | | | | | Zoe | Ava |
+--------+---------+---------+---------+---------+---------+---------+
| 14 | | | | | Joe | Zoe |
+--------+---------+---------+---------+---------+---------+---------+
| 15 | Joe | Ava | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 16 | | | | | Tom | Bob |
+--------+---------+---------+---------+---------+---------+---------+
| 17 | Sam | Bob | | | | |
+--------+---------+---------+---------+---------+---------+---------+
| 18 | | | Bob | Joe | | |
+--------+---------+---------+---------+---------+---------+---------+
| 19 | Tom | Mia | | | | |
+--------+---------+---------+---------+---------+---------+---------+
[Finished in 165ms]
Altitudes
Alice and Bob took a journey to the mountains. They have been climbing
up and down for N days and came home extremely tired.
Alice only remembers that they started their journey at an altitude of
H1 meters and they finished their wandering at an alitude of H2
meters. Bob only remembers that every day they changed their altitude
by A, B, or C meters. If their altitude on the ith day was x,
then their altitude on day i + 1 can be x + A, x + B, or x + C.
Now, Bob wonders in how many ways they could complete their journey.
Two journeys are considered different if and only if there exist a day
when the altitude that Alice and Bob covered that day during the first
journey differs from the altitude Alice and Bob covered that day during
the second journey.
Bob asks Alice to tell her the number of ways to complete the journey.
Bob needs your help to solve this problem.
Input format
The first and only line contains 6 integers N, H1, H2, A, B, C that
represents the number of days Alice and Bob have been wandering,
altitude on which they started their journey, altitude on which they
finished their journey, and three possible altitude changes,
respectively.
Output format
Print the answer modulo 10**9 + 7.
Constraints
1 <= N <= 10**5
-10**9 <= H1, H2 <= 10**9
-10**9 <= A, B, C <= 10**9
Sample Input
2 0 0 1 0 -1
Sample Output
3
Explanation
There are only 3 possible journeys-- (0, 0), (1, -1), (-1, 1).
Note
This problem comes originally from a hackerearth competition, now closed. The explanation for the sample input and output has been corrected.
Here is my solution in Python 3.
The question can be simplified from its 6 input parameters to only 4 parameters. There is no need for the beginning and ending altitudes--the difference of the two is enough. Also, we can change the daily altitude changes A, B, and C and get the same answer if we make a corresponding change to the total altitude change. For example, if we add 1 to each of A, B, and C, we could add N to the altitude change: 1 additional meter each day over N days means N additional meters total. We can "normalize" our daily altitude changes by sorting them so A is the smallest, then subtract A from each of the altitude changes and subtract N * A from the total altitude change. This means we now need to add a bunch of 0's and two other values (let's call them D and E). D is not larger than E.
We now have an easier problem: take N values, each of which is 0, D, or E, so they sum to a particular total (let's say H). This is the same at using up to N numbers equaling D or E, with the rest zeros.
We can use mathematics, in particular Bezout's identity, to see if this is possible. Some more mathematics can find all the ways of doing this. Once we know how many 0's, D's, and E's, we can use multinomial coefficients to find how many ways these values can be rearranged. Total all these up and we have the answer.
This code finds the total number of ways to complete the journey, and takes it modulo 10**9 + 7 only at the very end. This is possible since Python uses large integers. The largest result I found in my testing is for the input values 100000 0 100000 0 1 2 which results in a number with 47,710 digits before taking the modulus. This takes a little over 8 seconds on my machine.
This code is a little longer than necessary, since I made some of the routines more general than necessary for this problem. I did this so I can use them in other problems. I used many comments for clarity.
# Combinatorial routines -----------------------------------------------
def comb(n, k):
"""Compute the number of ways to choose k elements out of a pile of
n, ignoring the order of the elements. This is also called
combinations, or the binomial coefficient of n over k.
"""
if k < 0 or k > n:
return 0
result = 1
for i in range(min(k, n - k)):
result = result * (n - i) // (i + 1)
return result
def multcoeff(*args):
"""Return the multinomial coefficient
(n1 + n2 + ...)! / n1! / n2! / ..."""
if not args: # no parameters
return 1
# Find and store the index of the largest parameter so we can skip
# it (for efficiency)
skipndx = args.index(max(args))
newargs = args[:skipndx] + args[skipndx + 1:]
result = 1
num = args[skipndx] + 1 # a factor in the numerator
for n in newargs:
for den in range(1, n + 1): # a factor in the denominator
result = result * num // den
num += 1
return result
def new_multcoeff(prev_multcoeff, x, y, z, ag, bg):
"""Given a multinomial coefficient prev_multcoeff =
multcoeff(x-bg, y+ag, z+(bg-ag)), calculate multcoeff(x, y, z)).
NOTES: 1. This uses bg multiplications and bg divisions,
faster than doing multcoeff from scratch.
"""
result = prev_multcoeff
for d in range(1, ag + 1):
result *= y + d
for d in range(1, bg - ag + 1):
result *= z + d
for d in range(bg):
result //= x - d
return result
# Number theory routines -----------------------------------------------
def bezout(a, b):
"""For integers a and b, find an integral solution to
a*x + b*y = gcd(a, b).
RETURNS: (x, y, gcd)
NOTES: 1. This routine uses the convergents of the continued
fraction expansion of b / a, so it will be slightly
faster if a <= b, i.e. the parameters are sorted.
2. This routine ensures the gcd is nonnegative.
3. If a and/or b is zero, the corresponding x or y
will also be zero.
4. This routine is named after Bezout's identity, which
guarantees the existences of the solution x, y.
"""
if not a:
return (0, (b > 0) - (b < 0), abs(b)) # 2nd is sign(b)
p1, p = 0, 1 # numerators of the two previous convergents
q1, q = 1, 0 # denominators of the two previous convergents
negate_y = True # flag if negate y=q (True) or x=p (False)
quotient, remainder = divmod(b, a)
while remainder:
b, a = a, remainder
p, p1 = p * quotient + p1, p
q, q1 = q * quotient + q1, q
negate_y = not negate_y
quotient, remainder = divmod(b, a)
if a < 0:
p, q, a = -p, -q, -a # ensure the gcd is nonnegative
return (p, -q, a) if negate_y else (-p, q, a)
def byzantine_bball(a, b, s):
"""For nonnegative integers a, b, s, return information about
integer solutions x, y to a*x + b*y = s. This is
equivalent to finding a multiset containing only a and b that
sums to s. The name comes from getting a given basketball score
given scores for shots and free throws in a hypothetical game of
"byzantine basketball."
RETURNS: None if there is no solution, or an 8-tuple containing
x the smallest possible nonnegative integer value of
x.
y the value of y corresponding to the smallest
possible integral value of x. If this is negative,
there is no solution for nonnegative x, y.
g the greatest common divisor (gcd) of a, b.
u the found solution to a*u + b*v = g
v " "
ag a // g, or zero if g=0
bg b // g, or zero if g=0
sg s // g, or zero if g=0
NOTES: 1. If a and b are not both zero and one solution x, y is
returned, then all integer solutions are given by
x + t * bg, y - t * ag for any integer t.
2. This routine is slightly optimized for a <= b. In that
case, the solution returned also has the smallest sum
x + y among positive integer solutions.
"""
# Handle edge cases of zero parameter(s).
if 0 == a == b: # the only score possible from 0, 0 is 0
return (0, 0, 0, 0, 0, 0, 0, 0) if s == 0 else None
if a == 0:
sb = s // b
return (0, sb, b, 0, 1, 0, 1, sb) if s % b == 0 else None
if b == 0:
sa = s // a
return (sa, 0, a, 1, 0, 1, 0, sa) if s % a == 0 else None
# Find if the score is possible, ignoring the signs of x and y.
u, v, g = bezout(a, b)
if s % g:
return None # only multiples of the gcd are possible scores
# Find one way to get the score, ignoring the signs of x and y.
ag, bg, sg = a // g, b // g, s // g # we now have ag*u + bg*v = 1
x, y = sg * u, sg * v # we now have a*x + b*y = s
# Find the solution where x is nonnegative and as small as possible.
t = x // bg # Python rounds toward minus infinity--what we want
x, y = x - t * bg, y + t * ag
# Return the information
return (x, y, g, u, v, ag, bg, sg)
# Routines for this puzzle ---------------------------------------------
def altitude_reduced(n, h, d, e):
"""Return the number of distinct n-tuples containing only the
values 0, d, and e that sum to h. Assume that all these
numbers are integers and that 0 <= d <= e.
"""
# Handle some impossible special cases
if n < 0 or h < 0:
return 0
# Handle some other simple cases with zero values
if n == 0:
return 0 if h else 1
if 0 == d == e: # all step values are zero
return 0 if h else 1
if 0 == d or d == e: # e is the only non-zero step value
# If possible, return # of tuples with proper # of e's, the rest 0's
return 0 if h % e else comb(n, h // e)
# Handle the main case 0 < d < e
# --Try to get the solution with the fewest possible non-zero days:
# x d's and y e's and the rest zeros: all solutions are given by
# x + t * bg, y - t * ag
solutions_info = byzantine_bball(d, e, h)
if not solutions_info:
return 0 # no way at all to get h from d, e
x, y, _, _, _, ag, bg, _ = solutions_info
# --Loop over all solutions with nonnegative x, y, small enough x + y
result = 0
while y >= 0 and x + y <= n: # at most n non-zero days
# Find multcoeff(x, y, n - x - y), in a faster way
if result == 0: # 1st time through loop: no prev coeff available
amultcoeff = multcoeff(x, y, n - x - y)
else: # use previous multinomial coefficient
amultcoeff = new_multcoeff(amultcoeff, x, y, n - x - y, ag, bg)
result += amultcoeff
x, y = x + bg, y - ag # x+y increases by bg-ag >= 0
return result
def altitudes(input_str=None):
# Get the input
if input_str is None:
input_str = input('Numbers N H1 H2 A B C? ')
# input_str = '100000 0 100000 0 1 2' # replace with prev line for input
n, h1, h2, a, b, c = map(int, input_str.strip().split())
# Reduce the number of parameters by normalizing the values
h_diff = h2 - h1 # net altitude change
a, b, c = sorted((a, b, c)) # a is now the smallest
h, d, e = h_diff - n * a, b - a, c - a # reduce a to zero
# Solve the reduced problem
print(altitude_reduced(n, h, d, e) % (10**9 + 7))
if __name__ == '__main__':
altitudes()
Here are some of my test routines for the main problem. These are suitable for pytest.
# Testing, some with pytest ---------------------------------------------------
import itertools # for testing
import collections # for testing
def brute(n, h, d, e):
"""Do alt_reduced with brute force."""
return sum(1 for v in itertools.product({0, d, e}, repeat=n)
if sum(v) == h)
def brute_count(n, d, e):
"""Count achieved heights with brute force."""
if n < 0:
return collections.Counter()
return collections.Counter(
sum(v) for v in itertools.product({0, d, e}, repeat=n)
)
def test_impossible():
assert altitude_reduced(0, 6, 1, 2) == 0
assert altitude_reduced(-1, 6, 1, 2) == 0
assert altitude_reduced(3, -1, 1, 2) == 0
def test_simple():
assert altitude_reduced(1, 0, 0, 0) == 1
assert altitude_reduced(1, 1, 0, 0) == 0
assert altitude_reduced(1, -1, 0, 0) == 0
assert altitude_reduced(1, 1, 0, 1) == 1
assert altitude_reduced(1, 1, 1, 1) == 1
assert altitude_reduced(1, 2, 0, 1) == 0
assert altitude_reduced(1, 2, 1, 1) == 0
assert altitude_reduced(2, 4, 0, 3) == 0
assert altitude_reduced(2, 4, 3, 3) == 0
assert altitude_reduced(2, 4, 0, 2) == 1
assert altitude_reduced(2, 4, 2, 2) == 1
assert altitude_reduced(3, 4, 0, 2) == 3
assert altitude_reduced(3, 4, 2, 2) == 3
assert altitude_reduced(4, 4, 0, 2) == 6
assert altitude_reduced(4, 4, 2, 2) == 6
assert altitude_reduced(2, 6, 0, 2) == 0
assert altitude_reduced(2, 6, 2, 2) == 0
def test_main():
N = 12
maxcnt = 0
for n in range(-1, N):
for d in range(N): # must have 0 <= d
for e in range(d, N): # must have d <= e
counts = brute_count(n, d, e)
for h, cnt in counts.items():
if cnt == 25653:
print(n, h, d, e, cnt)
maxcnt = max(maxcnt, cnt)
assert cnt == altitude_reduced(n, h, d, e)
print(maxcnt) # got 25653 for N = 12, (n, h, d, e) = (11, 11, 1, 2) etc.
I'm trying to get to the result on what would be a double for-loop in another language (Java or JavaScript, for instance).
So the closest I can come up with is something like this:
1> L = [1,2,3].
[1,2,3]
2> R = [X + Y || X <- L, Y <- L].
[2,3,4,3,4,5,4,5,6]
3>
...but what I do really want is: [3,4,5]. I don't want to sum the elements that were already added:
A1 + A2
A2 + A3
A2 + A1 [already computed, position switched]
A2 + A3 [already computed, position switched]
A3 + A1
A3 + A2 [already computed, position switched]
Thanks in advance...
TL;DR
[X+Y || X <- L, Y <- L, Y > X].
Other solutions
You essentially want two iterators walking alongside the same data structure and an accumulator to collect sums of distinctive elements. There is no reason why you wouldn't be able to mimic such iterators in Erlang:
-module(sum2).
-export([start/1]).
start(Max) ->
L = lists:seq(1, Max),
T = list_to_tuple(L),
do_sum(T, 1, 2, size(T), []).
do_sum(T, X, S, S, A) when X + 1 =:= S ->
lists:reverse([mk_sum(X, S, T) | A]);
do_sum(T, X, S, S, A) ->
do_sum(T, X + 1, X + 2, S, [mk_sum(X, S, T) | A]);
do_sum(T, X, Y, S, A) ->
do_sum(T, X, Y + 1, S, [mk_sum(X, Y, T) | A]).
mk_sum(X, Y, T) -> element(X, T) + element(Y, T).
The result:
7> c(sum2).
{ok,sum2}
8> sum2:start(3).
[3,4,5]
9> sum2:start(5).
[3,4,5,6,5,6,7,7,8,9]
There is actually a simpler solution if you don't have a list of elements that you want to sum but just integers:
-module(sum3).
-export([start/1]).
start(Max) -> do_sum(1, 2, Max, []).
do_sum(X, S, S, A) when X + 1 =:= S -> lists:reverse([X + S | A]);
do_sum(X, S, S, A) -> do_sum(X + 1, X + 2, S, [X + S | A]);
do_sum(X, Y, S, A) -> do_sum(X, Y + 1, S, [X + Y | A]).
Or even a simpler solution with just list comprehension:
4> L = [1, 2, 3].
[1,2,3]
5> [X+Y || X <- L, Y <- L, Y > X].
[3,4,5]
6> f().
ok
7> L = [1,2,3,4,5].
[1,2,3,4,5]
8> [X+Y || X <- L, Y <- L, Y > X].
[3,4,5,6,5,6,7,7,8,9]
Also check this question, Erlang; list comprehension without duplicates, which tackles a similar problem and has more ideas for possible solutions.