class Vector: def __init__(self, x, y): self.x = x self.y = y def __add__(self, v2): return Vector(self.x + v2.x, self.y + v2.y) def __sub__(self, v2): return Vector(self.x - v2.x, self.y - v2.y) def __mul__(self, v2): if isinstance(v2, int): return Vector(self.x * v2, self.y * v2) else: return Vector(self.x * v2.x, self.y * v2.y) __rmul__ = __mul__ def __str__(self): return "{" + str(self.x) + ", " + str(self.y) + "}" def __repr__(self): return "{" + str(self.x) + ", " + str(self.y) + "}" def mag(self): # Magnitude squared. For our scenario, this is # irrelvant, since sqrt is a monotonic function. # Caches magnitude, need to update manually. self.magnitude = self.x**2 + self.y**2 return self def whole(self): # Gets reduced form of vector for parallel comparison. # Caches, need to update manually if self.x == 0 and self.y == 0: self.whole = (0,0) elif self.x == 0: self.whole = (0, sgn(self.y)) elif self.y == 0: self.whole = (sgn(self.x), 0) else: d = gcd(abs(self.x), abs(self.y)) self.whole = (self.x/d, self.y/d) return self def abs(a): return a if a > 0 else -a def sgn(a): return -1 if a < 0 else 1 def gcd(a, b): while b > 0: a, b = b, a % b return a def gen_spiral(w): x = [] i, j, dc = 0, 0, [0, -1] for z in range((2*w + 1)**2): x.append((i, j)) if i == j or (i < 0 and i == -j) or (i > 0 and i == 1 - j): dc = [-dc[1], dc[0]] i, j = i+dc[0], j+dc[1] return x def solution(dim, pos_y, pos_g, dist): dim, pos_y, pos_g = Vector(*dim), Vector(*pos_y), Vector(*pos_g) # Vectors of all mirrors of you and guard. count = 0 all_you_set, all_guard_set = set(), set() w = (dist / min([dim.x, dim.y])) + 1 dist = dist**2 spiral = gen_spiral(w) hits = set() off_y, off_g = (-2*pos_y + dim), (-2*pos_g + dim) for c in spiral: base = Vector(*c) * dim # Handle mirror reflections every other. decide = Vector(int(c[0] % 2 == 1), int(c[1] % 2 == 1)) mir_y = (base + decide * off_y).mag().whole() mir_g = ((base + pos_g + decide * off_g) - pos_y).mag().whole() # Check that new vectors don't already exist. y_len, g_len = len(all_you_set), len(all_guard_set) all_you_set.add(mir_y.whole) all_guard_set.add(mir_g.whole) uniq_y = y_len != len(all_you_set) uniq_g = g_len != len(all_guard_set) # Since we spiral out, any guard vector in set of you vectors must # hit you first. if mir_g.whole in all_you_set: continue if mir_g.magnitude <= dist and uniq_g: count += 1 return count