python_library_examples/PuLP/magic_square.py at main · krother/python_library_examples · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

from pulp import LpProblem, LpMinimize, LpVariable, LpInteger, lpSum, LpStatus, value

prob = LpProblem("PULPTEST", LpMinimize)

# model variables
XCOORD = [0, 1, 2]
YCOORD = [0, 1, 2]
NUMBERS = [1, 2, 3, 4, 5, 6, 7, 8, 9]

# variable is a 3 x 3 x 9 matrix of binary values
allocation = LpVariable.dicts("square", (XCOORD, YCOORD, NUMBERS), 0, 1, LpInteger)

# target function
prob += 0, "Arbitrary Objective Function"

# constraint: sum over rows
for x in XCOORD:
    prob += lpSum([n * allocation[x][y][n] for y in YCOORD for n in NUMBERS]) == 15

# constraint: sum over columns
for y in YCOORD:
    prob += lpSum([n * allocation[x][y][n] for x in XCOORD for n in NUMBERS]) == 15

# constraint: each number only once
for n in NUMBERS:
    prob += lpSum([allocation[x][y][n] for x in XCOORD for y in YCOORD]) == 1

# constraint: three numbers per column
for x in XCOORD:
    prob += lpSum([allocation[x][y][n] for y in YCOORD for n in NUMBERS]) == 3

# constraint: three numbers per row
for y in YCOORD:
    prob += lpSum([allocation[x][y][n] for x in XCOORD for n in NUMBERS]) == 3

# constraint: 9 numbers set
prob += lpSum([allocation[x][y][n] for x in XCOORD for y in YCOORD for n in NUMBERS]) == 9

# run the solver
prob.solve()

print("Status:", LpStatus[prob.status])

# print the numbers that have been found
for y in YCOORD:
    for x in XCOORD:
        for n in NUMBERS:
            if value(allocation[x][y][n]) == 1:
                print(n, end=' ')
    print()