from collections import defaultdict


def dict_with_lengths(words):
    """
    The function takes a list of words as its argument.
    Returns a dict with a length as key and a list of words with that length as value.

    For example, 
    INPUT: ["apple", "ball", "cat", "dog", "egg", "fruit"]
    OUTPUT: {3: ["ball", "cat", "dog"], 4: ["ball"], 5: ["apple", "fruit"]}
    """
    """
    This solution can be gradually improved from the original solution you had.

    # SOLUTION 1
    answer = dict()
    for word in words:
        key = len(word)
        key_words = answer.get(key, [])
        key_words.append(word)
        answer[key] = key_words

    # SOLUTION 2
    answer = dict()
    for word in words:
        key = len(word)
        answer.setdefault(key, [])
        answer[key].append(word)
    """
    # SOLUTION 3

    answer = defaultdict(list)
    for word in words:
        key = len(word)
        answer[key].append(word)

    return answer