algorithms/algorithms/matrix/bomb_enemy.py at main · keon/algorithms · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
"""
Bomb Enemy

Given a 2D grid, each cell is either a wall 'W', an enemy 'E' or empty '0'
(the number zero). Return the maximum enemies you can kill using one bomb.
The bomb kills all the enemies in the same row and column from the planted
point until it hits the wall since it is too strong to be destroyed.
You can only place the bomb at an empty cell.

Reference: https://leetcode.com/problems/bomb-enemy/

Complexity:
    Time:  O(m * n)
    Space: O(n)
"""

from __future__ import annotations


def max_killed_enemies(grid: list[list[str]]) -> int:
    """Return the maximum enemies killed by placing one bomb.

    Args:
        grid: 2D grid of 'W', 'E', and '0' characters.

    Returns:
        Maximum number of enemies killed.

    Examples:
        >>> max_killed_enemies([["0","E","0","0"],["E","0","W","E"],["0","E","0","0"]])
        3
    """
    if not grid:
        return 0
    rows, cols = len(grid), len(grid[0])
    max_killed = 0
    row_enemies, col_enemies = 0, [0] * cols
    for i in range(rows):
        for j in range(cols):
            if j == 0 or grid[i][j - 1] == "W":
                row_enemies = _row_kills(grid, i, j)
            if i == 0 or grid[i - 1][j] == "W":
                col_enemies[j] = _col_kills(grid, i, j)
            if grid[i][j] == "0":
                max_killed = max(max_killed, row_enemies + col_enemies[j])
    return max_killed


def _row_kills(grid: list[list[str]], row: int, col: int) -> int:
    """Count enemies killed in a row starting from the given column.

    Args:
        grid: 2D grid of characters.
        row: Row index.
        col: Starting column index.

    Returns:
        Number of enemies in the row segment.
    """
    count = 0
    num_cols = len(grid[0])
    while col < num_cols and grid[row][col] != "W":
        if grid[row][col] == "E":
            count += 1
        col += 1
    return count


def _col_kills(grid: list[list[str]], row: int, col: int) -> int:
    """Count enemies killed in a column starting from the given row.

    Args:
        grid: 2D grid of characters.
        row: Starting row index.
        col: Column index.

    Returns:
        Number of enemies in the column segment.
    """
    count = 0
    num_rows = len(grid)
    while row < num_rows and grid[row][col] != "W":
        if grid[row][col] == "E":
            count += 1
        row += 1
    return count