diff --git a/app/main.py b/app/main.py index 626f41cf..267ea138 100644 --- a/app/main.py +++ b/app/main.py @@ -1,34 +1,74 @@ class Deck: - def __init__(self, row, column, is_alive=True): - pass + def __init__(self, row: int, column: int, is_alive: bool = True) -> None: + self.row = row + self.column = column + self.is_alive = is_alive class Ship: - def __init__(self, start, end, is_drowned=False): + def __init__(self, start: tuple, + end: tuple, + is_drowned: bool = False) -> None: # Create decks and save them to a list `self.decks` - pass + self.is_drowned = is_drowned + self.decks = [] - def get_deck(self, row, column): + r1, c1 = start + r2, c2 = end + + if r1 == r2: + for col in range(min(c1, c2), max(c1, c2) + 1): + self.decks.append(Deck(r1, col)) + + elif c1 == c2: + for row in range(min(r1, r2), max(r1, r2) + 1): + self.decks.append(Deck(row, c1)) + + else: + raise ValueError("Ship must be horizontal or vertical") + + def get_deck(self, row: int, column: int) -> Deck | None: # Find the corresponding deck in the list - pass + for deck in self.decks: + if deck.row == row and deck.column == column: + return deck + return None - def fire(self, row, column): + def fire(self, row: int, column: int) -> str: # Change the `is_alive` status of the deck # And update the `is_drowned` value if it's needed - pass + deck = self.get_deck(row, column) + if deck is None or not deck.is_alive: + return "Miss!" + + deck.is_alive = False + if all(not d.is_alive for d in self.decks): + self.is_drowned = True + return "Sunk!" + + return "Hit!" class Battleship: - def __init__(self, ships): + def __init__(self, ships: list[Ship]) -> None: # Create a dict `self.field`. # Its keys are tuples - the coordinates of the non-empty cells, # A value for each cell is a reference to the ship # which is located in it - pass + self.field = {} + + for start, end in ships: + ship = Ship(start, end) + for deck in ship.decks: + self.field[(deck.row, deck.column)] = ship - def fire(self, location: tuple): + def fire(self, location: tuple) -> str: # This function should check whether the location # is a key in the `self.field` # If it is, then it should check if this cell is the last alive # in the ship or not. - pass + if location not in self.field: + return "Miss!" + ship = self.field[location] + row, column = location + return ship.fire(row, column)