Count colorings of a row of houses with fixed end colors

Clarify first. There are n houses in a row and k colors, numbered 0..k-1. Neighbors must have different colors. House 0 must be first and house n-1 must be last. Count the valid colorings modulo 10^9 + 7. Ask about these edge cases: with n == 1, the one house must be both first and last. A color outside 0..k-1 gives 0. With k == 1, only n == 1 works.

Brute force: backtracking. Fix house 0. Try every color for each next house that differs from its left neighbor. At the end, keep only the colorings that end on last. This takes O((k-1)^(n-1)) time. It is useful only to check small cases.

Key observation: only "same as first" matters. House 0 is fixed to first. Every other color plays the same role relative to it: swapping two non-first colors everywhere maps valid colorings to valid colorings. So after any number of houses, every color other than first has the same count. You only need to track whether house i equals first:

• same[i] = number of valid colorings of houses 0..i where house i has color first.
• diff[i] = number where house i has one particular color that is not first. By symmetry, every such color has the same count.

Transitions for the next house:

• To end on first, the previous house must be any of the k-1 other colors: same[i+1] = (k-1) * diff[i].
• To end on a particular other color c, the previous house can be first or any of the k-2 colors that are not first and not c: diff[i+1] = same[i] + (k-2) * diff[i].

Start with same[0] = 1, diff[0] = 0. The answer is same[n-1] if last == first, and diff[n-1] otherwise. This takes O(n) time and O(1) space. Common mistake: anchoring the symmetry on last instead. At house 0 only one color has count 1, so the non-last colors are not symmetric there, and the DP gives wrong answers (for example n=2, first=0, last=1).

Invariant. After step i, same + (k-1) * diff equals (k-1)^i, which is the total number of valid colorings of houses 0..i with no constraint on the last one. This is a quick sanity check in the interview.

Optimal for huge n: closed form. Let m = n - 1 be the number of steps. Using the invariant, same[i+1] = (k-1)*diff[i] = (k-1)^i - same[i]. Solving this recurrence (and using diff = ((k-1)^m - same) / (k-1)) gives:

• first == last: ((k-1)^m + (k-1) * (-1)^m) / k
• first != last: ((k-1)^m - (-1)^m) / k

The division is exact over the integers. Under a prime modulus, use the modular inverse of k (with pow(k, MOD - 2, MOD), valid when k is not a multiple of MOD). With fast exponentiation, this takes O(log n) time. Another way to see it: this is the number of walks of length m between two vertices of the complete graph K_k, which is the eigenvalue formula for K_k (eigenvalues k-1 once and -1 repeated k-1 times).

Present the DP first, since it is easy to explain. Then derive the closed form, or use 2x2 matrix exponentiation on (same, diff), when the interviewer pushes n to 10^18.

MOD = 1_000_000_007


def _validate(n: int, k: int, first: int, last: int) -> bool:
    if n < 1 or k < 1:
        raise ValueError("n and k must be positive")
    return 0 <= first < k and 0 <= last < k


def count_ways_dp(n: int, k: int, first: int, last: int) -> int:
    """O(n) time, O(1) space. Tracks only 'equals first' vs. one other color."""
    if not _validate(n, k, first, last):
        return 0
    same, diff = 1, 0                 # house 0 is `first`
    for _ in range(n - 1):
        same, diff = ((k - 1) * diff) % MOD, (same + (k - 2) * diff) % MOD
    return same if first == last else diff


def count_ways(n: int, k: int, first: int, last: int) -> int:
    """O(log n) closed form: walks of length n-1 in the complete graph K_k."""
    if not _validate(n, k, first, last):
        return 0
    m = n - 1
    sign = 1 if m % 2 == 0 else MOD - 1          # (-1)^m mod MOD
    power = pow(k - 1, m, MOD)                    # pow(0, 0) == 1, as needed
    if first == last:
        numerator = power + (k - 1) * sign
    else:
        numerator = power - sign
    return numerator % MOD * pow(k, MOD - 2, MOD) % MOD


def count_ways_brute(n: int, k: int, first: int, last: int) -> int:
    """Backtracking reference for small inputs."""
    if not _validate(n, k, first, last):
        return 0

    def go(i: int, prev: int) -> int:
        if i == n:
            return 1 if prev == last else 0
        return sum(go(i + 1, c) for c in range(k) if c != prev)

    return go(1, first)


if __name__ == "__main__":
    assert count_ways(1, 3, 2, 2) == 1
    assert count_ways(1, 3, 0, 1) == 0
    assert count_ways(2, 3, 0, 0) == 0
    assert count_ways(2, 3, 0, 1) == 1
    assert count_ways(3, 3, 0, 0) == 2
    assert count_ways(3, 3, 0, 1) == 1
    assert count_ways(4, 3, 0, 0) == 2
    assert count_ways(4, 3, 0, 2) == 3
    assert count_ways(5, 2, 0, 0) == 1
    assert count_ways(4, 2, 0, 0) == 0
    assert count_ways(3, 1, 0, 0) == 0
    assert count_ways(3, 4, 0, 5) == 0
    for n in range(1, 8):
        for k in range(1, 5):
            for a in range(k):
                for b in range(k):
                    want = count_ways_brute(n, k, a, b)
                    assert count_ways_dp(n, k, a, b) == want, (n, k, a, b)
                    assert count_ways(n, k, a, b) == want, (n, k, a, b)
    big = count_ways(10**18, 10**9, 3, 7)
    assert 0 <= big < MOD
    print("all tests passed")

• count_ways_dp: O(n) time, O(1) space.
• count_ways (closed form): O(log n + log MOD) time for the two modular powers, O(1) space.
• count_ways_brute: O((k-1)^(n-1)) time, O(n) recursion depth. Use it for tests only.

• n=1, k=3, first=2, last=2 -> 1 (the single house is both ends).
• n=1, k=3, first=0, last=1 -> 0.
• n=2, k=3, first=0, last=0 -> 0 (neighbors would match).
• n=2, k=3, first=0, last=1 -> 1.
• n=3, k=3, first=0, last=0 -> 2 (0,1,0 and 0,2,0).
• n=3, k=3, first=0, last=1 -> 1 (0,2,1).
• n=4, k=3, first=0, last=0 -> 2; first=0, last=2 -> 3.
• n=5, k=2, first=0, last=0 -> 1 (it must alternate); n=4, k=2, first=0, last=0 -> 0 (parity).
• n=3, k=1 -> 0, since neighbors cannot differ.
• last=5 with k=4 -> 0 (color out of range).
• n=10^18, k=10^9 -> runs instantly with the closed form; check it against the DP on small n.
• Cross-check: for all n <= 7, k <= 4 and every (first, last), the DP and the closed form match the backtracking count.

• Some middle houses are already painted. Split the row at the fixed houses. Each segment between two fixed colors is the same subproblem, with its length and both end colors given. Multiply the segment counts. If two fixed houses are adjacent and share a color, the answer is 0. This takes O(f log n) for f fixed houses after sorting them.
• Circular street (house n-1 is next to house 0) and no fixed colors. The count is (k-1)^n + (-1)^n (k-1), the chromatic polynomial of a cycle. Derive it as k choices for house 0 times the number of paths that end on a color different from it: k * (k-1) * diff[n-1]. If house 0 is fixed, divide by k.
• Each house has a cost per color; minimize total cost with fixed ends. This is Paint House II with fixed ends. Use the DP over colors with the "best and second-best of the previous row" trick: O(nk) time. Force house 0 and house n-1 to their colors by giving every other color infinite cost.
• Some color pairs cannot be neighbors (a general allowed-pair graph). Symmetry no longer holds. Use a k x k adjacency matrix A and compute entry (first, last) of A^(n-1). That is O(k^3 log n) time with matrix exponentiation, or O(nk^2) with a plain DP when n is small.
• Return one coloring, or list all of them. For one coloring, walk greedily and keep enough steps to reach last. Choose a color for house i such that the remaining count from it to last is nonzero, which the closed form tells you in O(1). To list all colorings, use backtracking with that same pruning, so the work stays proportional to the output.

Mid-level (L4)

• Asks about n == 1, first == last, and the modulus before coding.
• Writes a correct O(nk) DP over (house, color), then reduces it to O(n) with the "equals first" versus "other color" state when given a hint.
• Tests small cases by hand, including k = 2 parity.

Senior (L5)

• Finds the two-state symmetry without help, explains why every color other than first has the same count (and why anchoring on last breaks), and writes the O(n), O(1)-space DP.
• When n grows to 10^18, moves to 2x2 matrix exponentiation or the closed form and handles the modular inverse correctly.
• Checks the fast version against brute force on small inputs.

Staff (L6+)

• Derives the closed form (from the recurrence or the eigenvalues of K_k) and connects it to the chromatic polynomial of paths and cycles.
• Generalizes cleanly: pre-painted houses become a product over segments, forbidden pairs become a k x k matrix power, and costs become min-plus DP. Explains the cost trade-offs of each.
• Calls out numeric details: exact division before reduction, the inverse existing only when k mod p != 0, and pow(0, 0) == 1 for the k = 1, n = 1 case.