Sort words by a custom alphabet

The order string gives each character a rank: its index in order. Two
words compare the way they would in a dictionary, but with ranks in place of
the usual character codes. Walk both words from the left. The first position
where they differ decides the result. If one word runs out first, the shorter
word (the prefix) comes first.

Brute force. Write a comparator that calls order.index(c) for each
character it looks at, and pass it to the sort through
functools.cmp_to_key. This is correct, but each index call scans the order
string, so one comparison costs O(L · A), where L is the word length and A is
the alphabet size. Across O(n log n) comparisons that adds up.

Optimal. Build a dictionary rank[c] = i once, in O(A). Then turn every
word into a tuple of ranks and sort with that tuple as the key. Python
compares tuples element by element, and a tuple that is a prefix of another
is smaller. That is exactly the rule in the problem, so the key needs no
special case for prefixes. Each word is converted once (the key is cached by
sorted), and every comparison is a plain tuple comparison.

Some interviewers ask for the comparator form on purpose. The code shows it
too (compare), and it is also the function you need to check whether a list
is already sorted in this order.

Invariants

• rank is a bijection between the characters of order and
  0 .. A-1. A repeated character in order is an input error.
• Every character in every word appears in order. If one does not, fail
  loudly instead of guessing (see follow-ups for other policies).
• The sort is stable, so equal words keep their input order. Duplicates
  stay in the output.

from functools import cmp_to_key


def build_rank(order: str) -> dict[str, int]:
    """Map each character to its position in the custom alphabet."""
    rank: dict[str, int] = {}
    for i, c in enumerate(order):
        if c in rank:
            raise ValueError(f"duplicate character {c!r} in order")
        rank[c] = i
    return rank


def word_key(word: str, rank: dict[str, int]) -> tuple[int, ...]:
    try:
        return tuple(rank[c] for c in word)
    except KeyError as e:
        raise ValueError(f"character {e.args[0]!r} in {word!r} is not in order") from None


def custom_sort(order: str, words: list[str]) -> list[str]:
    """Sort words by the custom alphabet. Tuple comparison handles the prefix rule."""
    rank = build_rank(order)
    return sorted(words, key=lambda w: word_key(w, rank))


def compare(a: str, b: str, rank: dict[str, int]) -> int:
    """Return -1, 0 or 1, like a classic comparator."""
    for x, y in zip(a, b):
        if x != y:
            return -1 if rank[x] < rank[y] else 1
    # One word is a prefix of the other (or they are equal): shorter first.
    return (len(a) > len(b)) - (len(a) < len(b))


def custom_sort_cmp(order: str, words: list[str]) -> list[str]:
    """Comparator version, for interviewers who ask for one."""
    rank = build_rank(order)
    for w in words:
        word_key(w, rank)  # validate every character up front
    return sorted(words, key=cmp_to_key(lambda a, b: compare(a, b, rank)))


def is_sorted(order: str, words: list[str]) -> bool:
    """Check adjacent pairs only: sortedness is transitive."""
    rank = build_rank(order)
    return all(compare(a, b, rank) <= 0 for a, b in zip(words, words[1:]))


if __name__ == "__main__":
    order = "zyxwvutsrqponmlkjihgfedcba"
    assert custom_sort(order, ["apple", "zoo", "banana"]) == ["zoo", "banana", "apple"]
    assert custom_sort("cba", ["ab", "a", "c", "ba"]) == ["c", "ba", "a", "ab"]
    assert custom_sort("ab", []) == []
    assert custom_sort("ab", ["", "b", "a", ""]) == ["", "", "a", "b"]
    words = ["kite", "kit", "ki", "ant"]
    alpha = "abcdefghijklmnopqrstuvwxyz"
    assert custom_sort(alpha, words) == sorted(words)
    assert custom_sort_cmp(order, ["apple", "zoo", "banana"]) == ["zoo", "banana", "apple"]
    assert is_sorted("cba", ["c", "ba", "a", "ab"])
    assert not is_sorted("cba", ["ab", "a"])
    for bad in (lambda: custom_sort("aa", ["a"]), lambda: custom_sort("ab", ["abc"])):
        try:
            bad()
            raise AssertionError("expected ValueError")
        except ValueError:
            pass
    print("ok")

Let A be the alphabet size, n the number of words, and L the maximum word
length.

• build_rank: O(A) time and space.
• custom_sort (key version): building the keys is O(n · L). The sort does
  O(n log n) comparisons, each O(L) in the worst case, so O(n · L · log n)
  time. Keys take O(n · L) extra space.
• custom_sort_cmp: the same O(n · L · log n) time, with O(n) extra space
  for the sort (no stored keys), but a Python-level function call per
  comparison, so it is slower by a constant factor.
• is_sorted: O(total characters) time, O(A) space.

• order = "zyx...a" (reversed alphabet), ["apple", "zoo", "banana"] →
  ["zoo", "banana", "apple"].
• order = "cba", ["ab", "a", "c", "ba"] → ["c", "ba", "a", "ab"]
  ("a" is a prefix of "ab", so it comes first).
• Normal alphabet, ["kite", "kit", "ki", "ant"] → ["ant", "ki", "kit", "kite"].
• Empty list → [].
• Empty strings: order = "ab", ["", "b", "a", ""] → ["", "", "a", "b"].
• Duplicate words stay: order = "ba", ["a", "b", "a"] → ["b", "a", "a"].
• Single word → the same single word.
• order = "aa" → ValueError (duplicate in the order).
• order = "ab", ["abc"] → ValueError ('c' is not in the order).
• is_sorted("cba", ["ab", "a"]) → False.

1. Characters missing from the order. Ask what the policy is. Common
   choices: put unknown characters after all known ones (rank A + ord(c),
   which keeps them in their natural order among themselves), or ignore them.
   Only word_key changes.
2. Case-insensitive or multi-level ordering. Compare by the custom rank
   of the lowercased character first, then break ties by case. Make the key a
   pair: (tuple of ranks of lowercased chars, tuple of case flags).
3. Verify instead of sort. "Is this list sorted in the alien order?"
   (is_sorted). Only adjacent pairs need checking, so it is O(total
   characters) and no sort is needed.
4. Recover the order from a sorted list. Given words already sorted in an
   unknown alphabet, derive the alphabet. Compare each adjacent pair; the
   first differing characters give an edge x → y. Topologically sort the
   graph (Kahn's algorithm). A cycle, or a longer word before its own prefix,
   means the input is invalid. This is the classic "alien dictionary".
5. Very large input that does not fit in memory. Sort chunks in memory
   with the same key, write them out as runs, then do a k-way merge with a
   heap keyed by word_key. For many short words over a small alphabet, an
   MSD radix sort or a trie walked in rank order gives O(total characters)
   time.

Mid-level (L4)

• Builds a character-to-rank map instead of calling index on every
  comparison.
• Gets the prefix rule right (the shorter word first) without a bug at the
  end of the loop.
• Writes a working comparator or key, and tests an empty list, empty
  strings, and a prefix pair.

Senior (L5)

• Notices that a tuple of ranks as the sort key gives the prefix rule for
  free, and explains why key-based sorting is cheaper than a comparator in
  Python.
• Asks about unknown characters, duplicate order characters, case and
  duplicates in the input, and picks a clear policy for each.
• States the O(n · L · log n) cost and the space trade-off of storing keys.

Staff (L6+)

• Moves easily to the inverse problem (deriving the order by topological
  sort) and to checking sortedness in linear time.
• Discusses scale: external merge sort, or radix sort and tries when the
  alphabet is small and the words are short.
• Connects the problem to real collation (locale rules, Unicode
  normalization, multi-level comparison) and explains why a single rank per
  character is not enough there.