1
+ import re
2
+ from itertools import count
3
+
4
+ with open ("input" ) as f :
5
+ inp = f .read ().strip ()
6
+
7
+ first , second = inp .split ("\n \n " )
8
+
9
+
10
+ class Group :
11
+ def __init__ (self , type_ , n , hp , dmg , atype , init , weaks , immunes ):
12
+ self .type_ = type_
13
+ self .n = n
14
+ self .hp = hp
15
+ self .dmg = dmg
16
+ self .atype = atype
17
+ self .init = init
18
+ self .weaks = weaks
19
+ self .immunes = immunes
20
+
21
+ def epower (self ):
22
+ return self .n * self .dmg
23
+
24
+
25
+ def parse (d , text , type_ ):
26
+ for i , line in enumerate (text .split ("\n " )[1 :]):
27
+ n , hp , dmg , init = map (int , re .findall ("\d+" , line ))
28
+ temp = re .findall ("weak to (.*?)(?=[);]|$)" , line )
29
+ weaks = temp [0 ].split (", " ) if temp else []
30
+ temp = re .findall ("immune to (.*?)(?=[);]|$)" , line )
31
+ immunes = temp [0 ].split (", " ) if temp else []
32
+ atype = re .findall ("(\w+)\s+damage" , line )[0 ]
33
+ d [next (c )] = Group (type_ , n , hp , dmg , atype , init , weaks , immunes )
34
+ return d
35
+
36
+
37
+ def calc_dmg (grp1 , grp2 ):
38
+ if grp1 .atype in grp2 .immunes :
39
+ return 0
40
+ elif grp1 .atype in grp2 .weaks :
41
+ return grp1 .epower () * 2
42
+ else :
43
+ return grp1 .epower ()
44
+
45
+
46
+ def get_targets (groups ):
47
+ targets = []
48
+ attackers = sorted ([(i , grp .epower (), grp .init ) for i , grp in groups .items ()], key = lambda x : (- x [1 ], - x [2 ]))
49
+ used = set ()
50
+ for i , _ , _ in attackers :
51
+ attacker = groups [i ]
52
+ dmgs = [(j , calc_dmg (attacker , target ), target .epower (), target .init ) for j , target in groups .items () if j not in used and target .type_ != attacker .type_ ]
53
+ dmgs = sorted (dmgs , key = lambda x : (- x [1 ], - x [2 ], - x [3 ]))
54
+ if dmgs :
55
+ j , dmg , _ , _ = dmgs [0 ]
56
+ if dmg > 0 :
57
+ targets .append ((i , j , attacker .init ))
58
+ used .add (j )
59
+ return targets
60
+
61
+
62
+ def attack (targets , groups ):
63
+ res = 0
64
+ targets = sorted (targets , key = lambda x : - x [2 ])
65
+ dead = set ()
66
+ for i , j , _ in targets :
67
+ if i in dead :
68
+ continue
69
+ attacker = groups [i ]
70
+ target = groups [j ]
71
+ dmg = calc_dmg (attacker , target )
72
+ d_units = dmg // target .hp
73
+ res += d_units
74
+ groups [j ].n -= d_units
75
+ if groups [j ].n <= 0 :
76
+ dead .add (j )
77
+ groups .pop (j )
78
+ return res
79
+
80
+
81
+ def get_types (groups ):
82
+ res = set ()
83
+ for grp in groups .values ():
84
+ res .add (grp .type_ )
85
+ return res
86
+
87
+
88
+ def get_winner (boost = 0 ):
89
+ global c
90
+ c = count ()
91
+ groups = parse ({}, first , "immune" )
92
+ groups = parse (groups , second , "infected" )
93
+ for i , grp in groups .items ():
94
+ if grp .type_ == "immune" :
95
+ grp .dmg += boost
96
+ while True :
97
+ targets = get_targets (groups )
98
+ res = attack (targets , groups )
99
+ if res == 0 :
100
+ break
101
+ return get_types (groups ), sum (grp .n for grp in groups .values ())
102
+
103
+ # Part 1
104
+ _ , units = get_winner ()
105
+ print (units )
106
+
107
+ # Part 2
108
+ low , high = 0 , 10000
109
+ while (low < high ):
110
+ mid = low + (high - low ) // 2
111
+ remaining , units = get_winner (boost = mid )
112
+ if len (remaining ) > 1 :
113
+ low = mid + 1
114
+ else :
115
+ winner = remaining .pop ()
116
+ if winner == "immune" :
117
+ high = mid
118
+ else :
119
+ low = mid + 1
120
+
121
+ _ , units = get_winner (boost = low )
122
+ print (units )
0 commit comments