Scamming the coding interview: Problem 009: Find word in the matrix

find_word_in_the_matrix.py

"""
Scamming the coding interview
"""

from collections import defaultdict


def pre_process_matrix(matrix):

    letter_to_position_map = defaultdict(list)

    for i, row in enumerate(matrix):
        for j, letter in enumerate(row):
            letter_to_position_map[letter].append((i, j))

    return letter_to_position_map


def check_left_to_right_for_the_word(word, starting_position, matrix):

    for i, letter in enumerate(word):

        if starting_position[0] + i > len(matrix[0]) - 1:
            return i == len(word)

        if not letter == matrix[starting_position[0] + i][starting_position[1]]:
            return False

    return True


def check_top_to_bottom_for_the_word(word, starting_position, matrix):

    for i, letter in enumerate(word):

        if starting_position[1] + i > len(matrix) - 1:
            return i == len(word)

        if not letter == matrix[starting_position[0]][starting_position[1] + i]:
            return False

    return True


def find_word(matrix, word, letter_to_position_map):

    if not word:
        return False

    first_letter = word[0]

    first_letter_positions = letter_to_position_map[first_letter]

    for position in first_letter_positions:
        found = check_left_to_right_for_the_word(word, position, matrix) or \
            check_top_to_bottom_for_the_word(word, position, matrix)

        if found:
            return True

    return False


if __name__ == '__main__':

    matrix = [['F', 'A', 'C', 'I'],
        ['O', 'B', 'Q', 'P'],
        ['A', 'N', 'O', 'B'],
        ['M', 'A', 'S', 'S']]

    letter_to_position_map = pre_process_matrix(matrix)

    is_present = find_word(matrix, "AP", letter_to_position_map)

    print(is_present)