diff --git a/src/stack/asteroid_collision/__init__.py b/src/stack/asteroid_collision/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/stack/asteroid_collision/solution.py b/src/stack/asteroid_collision/solution.py new file mode 100644 index 0000000..23d5fd3 --- /dev/null +++ b/src/stack/asteroid_collision/solution.py @@ -0,0 +1,20 @@ +class Solution: + def asteroidCollision(self, asteroids: list[int]) -> list[int]: + stack = [] + for asteroid in asteroids: + while asteroid: + try: + left_asteroid = stack.pop() + if asteroid > 0 or (asteroid < 0 and left_asteroid < 0): + stack.extend([left_asteroid, asteroid]) + asteroid = None + else: + if abs(asteroid) == abs(left_asteroid): + asteroid = None + elif abs(asteroid) < abs(left_asteroid): + stack.append(left_asteroid) + asteroid = None + except IndexError: + stack.append(asteroid) + asteroid = None + return stack diff --git a/tests/test_asteroid_collision.py b/tests/test_asteroid_collision.py new file mode 100644 index 0000000..5f2e4fc --- /dev/null +++ b/tests/test_asteroid_collision.py @@ -0,0 +1,18 @@ +import pytest +from src.stack.asteroid_collision.solution import ( + Solution, +) + + +@pytest.mark.parametrize( + "asteroids, expected", + [ + ([5, 10, -5], [5, 10]), + ([8, -8], []), + ([10, 2, -5], [10]), + ([3, 5, -6, 2, -1, 4], [-6, 2, 4]), + ], +) +def test_asteroid_collision(asteroids, expected): + solution = Solution() + assert solution.asteroidCollision(asteroids) == expected