Bouncing Bullet Problem

I have inspected your code thoroughly and cannot find any issue; it seems entirely robust to me, at least in terms of correctness. However, I am aware that Foobar imposes time limits on problems; I suspect that your code failed that one test case because it was too slow.

To be precise, let us define $B$ as the maximum product of the sizes of allowed_x and allowed_y (in a vague sense, the maximum of the sizes of guard_hitting_bearings and self_hitting_bearings, although typically it will be much larger). We'll investigate the runtimes of each of the functions in your code, and provide a vague asymptotic bound.

Firstly, are_parallel is clearly $O(1)$, so we can just lump it in with the other constant time operations. The same goes for get_shorter_bearing.

Now, for get_all_bearings, we have an outer pair of loops that iterates through all $O(B)$ possibilities; and for each of the actually valid ones, it iterates through the entire list so far, for a total complexity of $O(B^2)$.

Similarly, for count_friendly_fires, we iterate over all pairs of bearings between friendly_bearings and guard_bearings, and thus again have a complexity of $O(B^2)$.

Overall, this code is $O(B^2)$ (plus some other small stuff we don't particularly care about), and we can significantly improve on that. Additionally, and importantly, this quadratic upper bound can be forced by a sufficiently bad input; for instance, setting your initial position to $(1, 1)$, the guard position to $(2, 2)$, and the room dimensions to $p \times p$ for some prime $p$ should force the sizes of guard_hitting_bearings and self_hitting_bearings to be at most a constant factor smaller than the product of the sizes of allowed_x and allowed_y. I suspect the last test case was probably something to this effect, designed to break solutions which were too slow (you probably passed the other test cases because they all had lots of parallel vectors, so the sizes reduced significantly). In any case, it is clear that improvements will have to be made to the two $O(B^2)$ functions listed above.

Looking at your code, a pattern emerges in the places that degenerate into $O(B^2)$: they all involve a naive parallel vector check, so optimising this is the key to making the code run faster. The key idea is as follows. Let $x \parallel y$ if and only if are_parallel(x, y). This is an equivalence relation ($x \parallel x$ for all $x$, and $x \parallel y$ and $y \parallel z$ implies $x \parallel z$). Now define $r(v)$ to be the vector with smallest magnitude such that $r(v) \parallel v$; it is easy to see that this vector is uniquely defined, and a function to compute $r(v)$ in $O(1)$ can be obtained by a simple modification of are_parallel. Now, since $r(v)$ is uniquely defined, we have the fact $$ x \parallel y \iff r(x) = r(y). $$ This is the key to optimising your solution. In particular, instead of arrays, we can use a dictionary indexed by $r(v)$.

Instead of the $O(B)$ inner loop in get_all_bearings to check if the new vector is parallel with anything already in bearing, we simply check the element at bearing[r(v)] (if such an element exists); this reduces the loop to a simple $O(1)$ dictionary lookup, and thus overall get_all_bearings is reduced to $O(B)$. Similarly, count_friendly_fires can iterate over all key-value pairs in friendly_bearings, and simply look up the element of guard_bearings parallel to the current key (if such an element exists); again the $O(B)$ inner loop is reduced to $O(1)$, making the function $O(B)$ overall. Thus, with this simple modification, your code can be made to run quite significantly faster.

Paul has some nice, readable code implementing this idea in another answer. His $r(v)$ is called get_direction, for reference.


I made my own implementation based on your idea and added code to compare the results between your and my implementation. They are the same. So your code is probably fine.

Here is the code, tested on Python 3.7:

import math

################################### Original solution ##########################################

def get_1D_bearings(t, m, g, distance):
    '''
    Returns a list of the total distance travelled in 1D which results in the beam arriving at the guard's t-coordinate.

    Parameters:
        t(int): The length of this dimension.
        m(int): The t-coodinate of the shooter (me).
        g(int): The t-coordinate of the guard.
        distance(int): The maximum allowed distance that the beam may travel.
    '''
    i = 0
    bearings = []
    path = g - m # Direct path from the shooter to the guard with no bounces.
    if abs(path) <= distance:
        bearings.append(path)
    while True:
        # Initial bounce off of the positive wall and final bounce off of the positive wall.
        path = (t - m) + (t-g) + 2*i*t
        if abs(path) <= distance:
            bearings.append(path)
        # Initial bounce off of the negative wall and final bounce off of the negative wall.
        path = - (m+g+2*i*t)
        if abs(path) <= distance:
            bearings.append(path)
        # Initial bounce off of the positive wall and final bounce off of the negative wall.
        path = (t - m) + g + (2*i+1)*t
        if abs(path) <= distance:
            bearings.append(path)
        # Initial bounce off of the negative wall and final bounce off of the positive wall.
        path = - ( m + (t - g) + (2*i+1)*t)
        if abs(path) <= distance:
            bearings.append(path)
        else:
            break
        i += 1
    return bearings

def are_parallel(a, b):
    '''
    Returns if the bearings given by a and b are parallel vectors.

    Parameters:
        a(array-like): A 2D-array of integers.
        b(array-like): A 2D-array of integers.
    '''
    x1, y1 = a
    x2, y2 = b
    div1 = abs(math.gcd(x1, y1))
    div2 = abs(math.gcd(x2, y2) )
    if div1 == 0 or div2 ==0:
        if not div1 == 0 and div2 == 0:
            return False
        elif (x1 == 0 and x2 == 0) and (y1 // abs(y1) == y2 // abs(y2)):
            return True
        elif (y1 == 0 and y2 ==0) and (x1 // abs(x1) == x2 // abs(x2)):
            return True
        else:
            return False
    else:
        if x1 // div1 == x2 // div2 and y1 // div1 == y2 // div2:
            return True
        else:
            return False

class VectorsNotParallel(Exception):
    '''Raise this exception when handling vectors which are assumed to be parallel but are not.'''
    def __init__(self):
        pass

def get_shorter_bearing(a, b):
    '''
    Returns the shorter vector of a and b. If they are not parallel, raises VectorsNotParallel.

    Parameters:
        a(array-like): A 2D-array of integers indicating a direction.
        b(array-like): A 2D-array of integers indicating a direction.
    '''
    if not are_parallel(a, b):
        raise VectorsNotParallel("These bearings point in different directions: " + str(a) + " and " + str(b))
    x1, y1 = a
    x2, y2 = b
    if x1 == 0:
        if abs(y1) < abs(y2):
            return a
        else:
            return b
    if y1 == 0:
        if abs(x1) < abs(x2):
            return a
        else:
            return b
    div1 = abs(math.gcd(x1, y1))
    div2 = abs(math.gcd(x2, y2) )
    if div1 < div2:
        return a
    else:
        return b

def get_all_bearings(dimensions, your_position, guard_position, distance):
    '''
    Combines the allowed distances from each of the two dimensions to generate a list of all of the
    allowed directions that can be shot in which take a beam from your_position to guard_position
    while not travelling further than the provided distance. Note that some of these directions include
    passing through your_position.

    Parameters:
        dimensions(array-like): A 2D-array of integers indicating the size of the room.
        your_position(array-like): A 2D-array of integers indicating your position in the room.
        guard_position(array-like): A 2D-array of integers indicating the guard's position in the room.
        distance(int): An integer indicating the maximum distance the beam can travel.

    Returns:
        bearings(array-like): An array of 2D-arrays indicating the bearings which move the beam from your_position
                              to guard_position.
    '''
    dx, dy=  dimensions
    sx, sy = your_position
    gx, gy = guard_position
    allowed_x = get_1D_bearings(dx, sx, gx, distance)
    allowed_y = get_1D_bearings(dy, sy, gy, distance)
    bearings = []
    for x in allowed_x:
        for y in allowed_y:
            if x**2 + y**2 < 1 or x**2 + y**2 > distance **2:
                continue
            res = [x, y]
            append = True  # Do we need to append to the list of bearings or just update an existing one
            for bearing in bearings:
                if are_parallel(res, bearing):
                    append = False
                    res_2 = get_shorter_bearing(res, bearing)
                    bearing[0] = res_2[0]
                    bearing[1] = res_2[1]
            if append:
                bearings.append(res)
    return bearings

def count_friendly_fires(friendly_bearings, guard_bearings):
    '''
    Returns the number of bearings which result in the guard being hit only after the beam
    passes through the shooter (which is not allowed).

    Parameters:
        friendly_bearings(array-like): An array of 2D arrays which indicate bearings that reach the shooter.
        guard_bearings(array-like): An array of 2D arrays which indicate bearings that reach the guard.
    '''
    count = 0
    for f_bearing in friendly_bearings:
        for g_bearing in guard_bearings:
            if are_parallel(f_bearing, g_bearing):
                if get_shorter_bearing(f_bearing, g_bearing) == f_bearing:
                    #print(f_bearing, g_bearing)
                    count += 1
    return count

def solution(dimensions, your_position, guard_position, distance):
    '''
    Returns the number of distinct directions that take a bullet from your_position to
    guard_position within the allowed distance.

    Parameters:
        dimensions(array-like): A 2D-array of integers indicating the size of the room.
        your_position(array-like): A 2D-array of integers indicating your position in the room.
        guard_position(array-like): A 2D-array of integers indicating the guard's position in the room.
        distance(int): An integer indicating the maximum distance the beam can travel.
    '''
    guard_hitting_bearings = get_all_bearings(dimensions, your_position, guard_position, distance)
    self_hitting_bearings = get_all_bearings(dimensions, your_position, your_position, distance)
    count = count_friendly_fires(self_hitting_bearings, guard_hitting_bearings)
    return len(guard_hitting_bearings) - count


################################### My solution ##########################################

def get_range (a, b, max_distance):
    """
    Gets values of a + bk for integers k such that -max_distance <= a + bk <= max_distance.
    :param a: integer
    :param b: positive integer
    :param max_distance: non-negative integer
    """
    return list(range((a + max_distance) % b - max_distance, max_distance + 1, b))

def get_1d_list (size, source, to, max_distance):
    return get_range(to - source, 2*size, max_distance) + get_range(-source - to, 2*size, max_distance)

def get_direction (dx, dy):
    gcd = math.gcd(dx, dy)
    return dx // gcd, dy // gcd

def get_direction_distance_map (dx_list, dy_list, max_squared_distance):
    """Returns a map which has directions as key and squared distances as value"""
    map = {}
    for dx in dx_list:
        for dy in dy_list:
            sd = dx*dx + dy*dy
            if 0 < sd <= max_squared_distance:
                dir = get_direction(dx, dy)
                if sd < map.get(dir, max_squared_distance + 1):
                    map[dir] = sd
    return map

def my_solution (dimensions, your_position, guard_position, distance):
    self_x = get_1d_list(dimensions[0], your_position[0], your_position[0], distance)
    self_y = get_1d_list(dimensions[1], your_position[1], your_position[1], distance)
    guard_x = get_1d_list(dimensions[0], your_position[0], guard_position[0], distance)
    guard_y = get_1d_list(dimensions[1], your_position[1], guard_position[1], distance)
    map_self = get_direction_distance_map(self_x, self_y, distance*distance)
    map_guard = get_direction_distance_map(guard_x, guard_y, distance*distance)

    # Remove friendly fires
    for dir, sd_self in map_self.items():
        sd_guard = map_guard.get(dir)
        if sd_guard is not None and sd_self < sd_guard:
            del map_guard[dir]

    return len(map_guard)


################################### Test code ##########################################

for w in range(2, 8):
    print(f"Testing width {w}")
    for d in range(16):
        for h in range(2, 8):
            for xs in range(1, w):
                for ys in range(1, h):
                    for xg in range(1, w):
                        for yg in range(1, h):
                            if xs != xg or ys != yg:
                                s1 = solution([w, h], [xs, ys], [xg, yg], d)
                                s2 = my_solution([w, h], [xs, ys], [xg, yg], d)
                                if s1 != s2:
                                    print(w, h, xs, ys, xg, yg)