leetcode algorithms 1320 - minimum_distance_to_type_a_word_using_two_fingers

1320_minimum_distance_to_type_a_word_using_two_fingers.py

import unittest
import string
import collections
from pprint import pprint

# def print(*args, **kwargs):
#     pass

def make_table(letters, width):
    table = {}
    i = 0
    l = 0
    while True:
        for j in range(width):
            ch = letters[i]
            table[ch] = (l, j)
            i += 1
            if i == len(letters):
                return table
        l += 1

def distance(table, word, l1, l2):
    if l1 < 0 or l2 < 0:
        return 0
    ch1_pos = table[word[l1]]
    ch2_pos = table[word[l2]]
    dis = abs(ch1_pos[0]-ch2_pos[0]) + abs(ch1_pos[1]-ch2_pos[1])
    return dis

class Solution:
    def __init__(self):
        self.counter = collections.defaultdict(int)
        self.cache = {}

    def minimumDistance(self, word: str) -> int:
        self.word = word
        self.table = make_table(string.ascii_uppercase, 6)
        l = len(word)
        return min([self.dp(i, l-1) for i in range(-1, l-1)]+ [self.dp(l-1, i) for i in range(-1, l-1)])
    
    def dp(self, l, r):
        if (l,r) in self.cache:
            return self.cache[(l,r)]
        self.counter[(l,r)] += 1
        if l == -1 and r >= 0:
            return sum(distance(self.table, self.word, i, j) for i, j in zip(range(r), range(1,r+1)))
        if l >= 0 and r == -1:
            return sum(distance(self.table, self.word, i, j) for i, j in zip(range(l), range(1,l+1)))
        if l == r:
            return self.dp(l-1, r)
        if l > r:
            move_left = self.dp(l-1, r) + distance(self.table, self.word, l-1, l)
            print(f"l > r, l:{l}, r:{r}, finger l on {self.word[l-1]}, finger r on {self.word[r]}, move finger 1 to {self.word[l]}, get cost {move_left}")
            if r == l-1:
                move_right = min((self.dp(pl, l-1) + distance(self.table, self.word, pl, l)) for pl in range(-1, l-1))
                dp_l_r = move_right
                print(f"l > r, l:{l}, r:{r}, iterate through finger l on -1 to {l-2}, finger r on {self.word[l-1]}, move finger 1 to {self.word[l]}, get cost {move_right}")
            else:
                dp_l_r = move_left
            print(f"dp[{l}][{r}]: {dp_l_r}")
            self.cache[(l,r)] = dp_l_r
            return dp_l_r
        if r > l:
            move_right = self.dp(l, r-1) + distance(self.table, self.word, r-1, r)
            print(f"r > l, l:{l}, r:{r}, finger r on {self.word[r-1]}, finger l on {self.word[l]}, if move finger r to {self.word[r]}, get cost {move_right}")
            if l == r - 1:
                move_left = min((self.dp(r-1, pr) + distance(self.table, self.word, pr, r)) for pr in range(-1, r-1))
                dp_l_r = move_left
                print(f"r > l, l:{l}, r:{r}, iterate through finger r on -1 to {r-2}, finger l on {self.word[r-1]}, if move finger r to {self.word[r]}, get cost {move_left}")
            else:
                dp_l_r = move_right
            print(f"dp[{l}][{r}]: {dp_l_r}")
            self.cache[(l,r)] = dp_l_r
            return dp_l_r

        
class T(unittest.TestCase):

    def setUp(self):
        table_width = 6
        self.table = make_table(string.ascii_uppercase, table_width)
    
    def test_minumum_distance_case1(self):
        self.s = Solution()
        word = "CAKE"
        self.assertEqual(3, self.s.minimumDistance(word))

    def test_minumum_distance_case2(self):
        self.s = Solution()
        word = "HAPPY"
        self.assertEqual(6, self.s.minimumDistance(word))

    def test_minumum_distance_case3(self):
        self.s = Solution()
        word = "NEW"
        self.assertEqual(3, self.s.minimumDistance(word))

    def test_minumum_distance_case4(self):
        self.s = Solution()
        word = "YEAR"
        # pprint(self.s.counter)
        self.assertEqual(7, self.s.minimumDistance(word))

    def test_minumum_distance_case5(self):
        self.s = Solution()
        word = "QIBZR"
        # pprint(self.s.counter)
        self.assertEqual(7, self.s.minimumDistance(word))

    def test_minumum_distance_case6(self):
        self.s = Solution()
        word = "OPVUWZLCKTDPSUKGHAXIDWHLZFKNBDZEWHBSURTVCADUGTSDMCLDBTAGFWDPGXZBVARNTDICHCUJLNFBQOBTDWMGILXPSFWVGYBZVFFKQIDTOVFAPVNSQJULMVIERWAOXCKXBRI"
        # pprint(self.s.counter)
        self.assertEqual(295, self.s.minimumDistance(word))

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