-
Notifications
You must be signed in to change notification settings - Fork 0
/
smt.py
80 lines (60 loc) · 2.8 KB
/
smt.py
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from common import parse_input, parse_output
from z3 import Optimize, Int, And, Implies, Sum, Bool, AtMost, Or, Not
# n: number of units
# k: number of periods
# areas: areas of the units
# neighbours: dim: n x neighbours
# profits: dim: k x n
# amin: min area of the natural reserve
def main():
import sys
n, k, areas, neighbours, profits, amin = parse_input(sys.stdin.read())
# print(f"{n=}\n{k=}\n{areas=}\n{neighbours=}\n{profits=}\n{amin=}") # debug
solver = Optimize()
periods = [[Bool(f"H{i}{j}") for j in range(1, k + 1)] for i in range(1, n + 1)]
nr = [Bool(f"NR{i}") for i in range(1, n + 1)]
depths = [Int(f'D{i}') for i in range(1, n + 1)]
for idx, (i_periods, nri) in enumerate(zip(periods, nr)):
# A unit of land can only be harvested once :
solver.add(AtMost(*i_periods, 1))
# Pi neighbours must not have the same value
for i_n in neighbours[idx]:
if idx + 1 > i_n:
for pj_i, pj_in in zip(i_periods, periods[i_n-1]):
solver.add(Or(Not(pj_i), Not(pj_in)))
# If natural reserve, cannot be harvested
if amin > 0:
solver.add(Implies(nri, And(*[Not(pj_i) for pj_i in i_periods])))
if amin > 0:
for idx, (nri, di) in enumerate(zip(nr, depths)):
solver.add(
# 0 <= di, # Di >= 0, Di = 0 : i does not belong to the tree, Di = d : i belongs to the tree a d level
Implies(1 <= di, nri), # Di >= 1 -> NR (Pi == k + 1)
Implies(nri, 1 <= di), # NR (Pi == k + 1) -> Di >= 1
Implies(di > 1, Sum(*[di == depths[i_n - 1] + 1 for i_n in neighbours[idx]]) == 1),
# i with depth d can only have 1 neighbour with depth d-1
)
# Must only have one root
solver.add(Sum(*[di == 1 for di in depths]) == 1)
# Area of natural reserve must be greater than Amin
solver.add(Sum(*[areas[idx] * nri for idx, nri in enumerate(nr)]) >= amin)
# Maximize profit
solver.maximize(
Sum(*[profits[j_idx][i_idx] * periods[i_idx][j_idx] for i_idx in range(n) for j_idx in range(k)]))
if not solver.check():
raise Exception("UNSAT")
sol = solver.model()
# Decoding
total_profit = 0
harvesting_periods: list[list[int]] = [[] for _ in range(k)]
natural_reserve = []
for i_idx, i_periods in enumerate(periods):
for j_idx, pj_i in enumerate(i_periods):
if sol[pj_i]:
harvesting_periods[j_idx].append(i_idx + 1)
total_profit += profits[j_idx][i_idx]
if sol[nr[i_idx]]:
natural_reserve.append(i_idx + 1)
print(parse_output(total_profit, harvesting_periods, natural_reserve), end="")
if __name__ == '__main__':
main()