"""Solve a programming puzzle: find the longest "word chain."

This was probably from Hackerrank or something.

The idea here is to find the *length of the* longest chain of words in
some language (here defined by a list of words supplied as an
argument) in which each word can be made from the previous word in the
chain by adding one letter.

"""


def longestChain(words):
    by_length = {}
    for word in words:
        if len(word) not in by_length:
            by_length[len(word)] = set()
        by_length[len(word)].add(word)
    
    max_chain_so_far = 0
    next_max_chains = {}
    for length in sorted(by_length.keys()):
        max_chains = next_max_chains
        next_max_chains = {}
        our_words = by_length[length]
        for word in our_words:
            if word not in max_chains:
                max_chains[word] = 1
            max_chain_so_far = max(max_chain_so_far, max_chains[word])

        # Iterate over the words in the next level (n < 50000) to see if
        # any of their one-char-removed versions (m < 60) are in the current level.
        # We could optimize this more, since either
        # the first half or second half of the word at the next level
        # needs to be present in some word at the current level.
        edges = [(our_word, next_word)
                 for next_word in by_length.get(length+1, [])
                 for our_word in [next_word[:i] + next_word[i+1:]
                                  for i in range(len(next_word))]
                 if our_word in our_words
                ]
        for our_word, next_word in edges:
            next_max_chains[next_word] = max(next_max_chains.get(next_word, 0),
                                             max_chains[our_word] + 1)
    
    return max_chain_so_far