battleship_solution.py

import unittest


def solution(N, S, T):
    start_ord = ord('A')
    
    s_arr = S.split(',')
    t_arr = T.split(' ')

    ships = []
    for s in s_arr:
        ship = {}

        split_ship = s.split(' ')

        start_coord = split_ship[0]
        end_coord = split_ship[1]

        start_x = ord(start_coord[1]) - start_ord
        start_y = ord(start_coord[0]) - ord('1')

        end_x = ord(end_coord[1]) - start_ord
        end_y = ord(end_coord[0]) - ord('1')

        size = (abs(start_x - end_x) + 1) * (abs(start_y - end_y) + 1)

        ship['start_x'] = start_x
        ship['start_y'] = start_y
        ship['end_x'] = end_x
        ship['end_y'] = end_y
        ship['size'] = size
    
        ships.append(ship)

    targets = []
    for t in t_arr:
        target = {}

        x = ord(t[1]) - start_ord
        y = ord(t[0]) - ord('1')

        target['x'] = x
        target['y'] = y

        targets.append(target)
    
    ships_hit = 0
    sank = 0
    
    for ship in ships:
        hits_on_ship = 0

        for target in targets:
            if is_hit_successful(target, ship):
                hits_on_ship += 1

        if hits_on_ship > 0:
            if hits_on_ship == ship['size']:
                sank += 1
            else:
                ships_hit += 1

    return '{0},{1}'.format(sank, ships_hit)   


def is_hit_successful(target, ship):
    is_hit = False

    check_upper_bounds = is_coords_greater(target['x'], target['y'], ship['start_x'], ship['start_y'])
    check_lower_bounds = is_coords_greater(ship['end_x'], ship['end_y'], target['x'], target['y'])

    if check_upper_bounds and check_lower_bounds:
        is_hit = True

    return is_hit


def is_coords_greater(x1, y1, x2, y2):
    return x1 >= x2 and y1 >= y2


class battleshiptest(unittest.TestCase):
    def test1(self):
        N = 4
        S = '1B 2C,2D 4D'
        T = '2B 2D 3D 4D 4A'
        expected = '1,1'
        self.assertEqual(solution(N, S, T), expected)

    def test2(self):
        N = 3
        S = '1A 1B,2C 2C'
        T = '1B'
        expected = '0,1'
        self.assertEqual(solution(N, S, T), expected)

    def test3(self):
        N = 12
        S = '1A 2A,12A 12A'
        T = '12A'
        expected = '1,0'
        self.assertEqual(solution(N, S, T), expected)


if __name__ == "__main__":
    unittest.main()