from typing import List

debug: bool = True


class Solution:
    def sumPrefixScores(self, words: List[str]) -> List[int]:
        word_count = len(words)
        # sort into dictionary order
        # this is the KEY
        sorted_indices = sorted(range(word_count), key=lambda i: words[i])
        common_prefix_lengths = self._calculate_common_prefix_lengths(
            words, sorted_indices
        )
        scores = self._calculate_scores(words, sorted_indices, common_prefix_lengths)
        return scores

    def _calculate_common_prefix_lengths(self, words, sorted_indices):
        common_prefix_lengths = [0] * len(words)
        for i in range(1, len(words)):
            prev_word = words[sorted_indices[i - 1]]
            curr_word = words[sorted_indices[i]]

            if debug:
                print(f"Prev: {prev_word},  Curr: {curr_word}")

            common_length = 0
            while (
                common_length < len(prev_word)
                and common_length < len(curr_word)
                and prev_word[common_length] == curr_word[common_length]
            ):
                common_length += 1

            if debug:
                print(f"CommonPrefixLength[<word: {words[i]}>]={common_length}")

            common_prefix_lengths[i] = common_length
        if debug:
            print("Common prefix lengths: ", common_prefix_lengths)
        return common_prefix_lengths

    def _calculate_scores(self, words, sorted_indices, common_prefix_lengths):
        scores = [0] * len(words)
        for i, word_index in enumerate(sorted_indices):
            word_length = len(words[word_index])
            scores[word_index] += word_length
            j = i + 1
            common_length = word_length

            # bidirectional update
            # - update scores of current word
            # - update scores of all following words, which's score changes caused by this word
            if debug:
                print(f"Bidirectional update using word: {words[word_index]}")
            while j < len(words):
                common_length = min(common_length, common_prefix_lengths[j])
                if common_length == 0:
                    break
                if debug:
                    print(
                        f"Update word: {words[sorted_indices[j]]}, delta: {common_length}"
                    )
                scores[word_index] += common_length
                scores[sorted_indices[j]] += common_length

                j += 1
        return scores


print(Solution().sumPrefixScores(["abc", "ab", "bc", "b"]))