IOAI Training: Learn AI by solving problems

The problem in brief

A friend thinks of a Swedish word from a fixed list of 10 001 words and gives you a starting word. After each of your guesses he answers with two words: whether your guess is "hotter" or "colder" (semantically closer to or farther from his word) than your previous guess, and than your best guess so far. You keep guessing until you hit the word. Your program plays $N$ rounds interactively (up to 1000 rounds in the last test group) within 8 seconds of CPU time. No training data is provided: you may bring any word vectors you like from the internet, but the program runs offline on Kattis, so whatever knowledge you use must fit inside the source file.

Important caveat for this correction. The judge's similarity model is secret ("he reads news and books"). We cannot reproduce it, so every score below is measured against a simulated judge that we built ourselves: cosine similarity between fastText vectors trained on Swedish Common Crawl. Our own program deliberately uses different vectors (fastText trained on Swedish Wikipedia), so that it faces a realistic model mismatch. Real Kattis scores may be higher or lower.

How you are scored

If a round needs $k$ guesses, it earns

\text{Score}(k) = \min\Big(100,\ 100\cdot\Big(\frac{10001-k}{10000}\Big)^{13}\Big).

The exponent 13 makes this a harsh curve: 10 guesses give 99.9 points, 100 guesses 87.7, 500 guesses 51.3, 1000 guesses 25.4 and 3000 guesses less than 1. A test case scores the average over its rounds, and the three groups ( $N \le 10$ , $N \le 100$ , $N \le 1000$ ) are worth 30, 40 and 30 points.

There is no published threshold. The provided baseline guesses the vocabulary in alphabetical order, so $k$ is uniform on $1..10001$ and its expected score is

100\int_0^1 (1-u)^{13}\,du = \frac{100}{14} \approx 7.1 .

First look at the data

The only attachment is vocabulary.txt: 10 001 lower-case Swedish words, mostly frequent news words (regeringen, miljoner), first names and place names, plus a surprising layer of EU-legislation jargon (associeringsavtal, budgetförordningen) and even a fungus class (agaricomycetes). This tells us the judge's corpus mixes news, books, Wikipedia and EU documents, and that some words will be missing from any vector set we pick: from the first 300 MB of the Wikipedia fastText file, 533 of the 10 001 words had no vector.

Two constraints shape everything:

Size. Kattis accepts at most 128 KiB of source code. The vocabulary alone takes 31 KB once compressed with lzma and encoded in base85. That leaves room for about 6 to 7 bytes per word, while a fastText vector is 300 floats (1200 bytes). We must compress by a factor of 200.
Time. With 1000 rounds in 8 s, a round may use about 7 ms of CPU, input and output included.

A first baseline

Alphabetical order scores 7.1 on average. A natural first idea is hill climbing: guess a neighbour of the best word so far and move whenever the answer is "hotter". It throws away most of the information: every answer is a statement about the hidden word, and all of them should be combined.

The key idea

Write $\mathbf{t}$ for the (unit) vector of the hidden word, and $\mathbf{g}$ , $\mathbf{b}$ for the guess and the word it is compared to. The answer " $g$ is hotter than $b$ " means

\cos(\mathbf{t},\mathbf{g}) > \cos(\mathbf{t},\mathbf{b}) \iff \mathbf{t}\cdot(\mathbf{g}-\mathbf{b}) > 0 .

Every answer is a half-space: a hyperplane through the origin, perpendicular to $\mathbf{g}-\mathbf{b}$ , cuts the space of words in two, and the target lives on the announced side. Each guess gives two such cuts (versus the previous guess and versus the best guess). If our vectors were exactly the judge's, each cut would remove about half of the remaining candidates and we would need about $\log_2 10001 / 2 \approx 7$ guesses after the first ones: in simulation, with the judge's own vectors, the median is 13 guesses and the score 98.

But our vectors are not the judge's, so some cuts will be "wrong" for the true target. Hard elimination would then throw the target away for ever. Instead we keep a soft score, the negative log-likelihood of each candidate $c$ under a noisy-answer model:

P(\text{answer} \mid c) = \sigma\big(\beta\, s\, \mathbf{c}\cdot(\mathbf{g}-\mathbf{b})\big),\qquad \text{pen}(c) \mathrel{+}= \log\big(1+e^{-\beta\, s\, \mathbf{c}\cdot(\mathbf{g}-\mathbf{b})}\big),

with $s=+1$ for "hotter", $s=-1$ for "colder", $\sigma$ the logistic function and $\beta = 20$ . A candidate barely on the wrong side of a cut pays little, one far on the wrong side pays a lot. We always guess the unguessed word with the smallest penalty (the most likely target). This is exactly logistic-regression-style Bayesian updating, and it is what makes the method robust to a mismatched model: on 150 simulated rounds the soft score beat hard "count the violated cuts" (92.9 vs 88.7 with full vectors).

The full solution

solution.py builds the submission; the generated program is about 118 KB.

Vectors. fastText Swedish Wikipedia vectors filtered to the vocabulary (lower-case match, falling back to a capitalised form for names). A word without a vector gets the mean vector of the known words sharing its longest prefix of at least 4 letters, since Swedish inflections share their stem (andades borrows from andas): this rescued 375 of the 533 missing words.
Compression: PCA + product quantisation. We centre the vectors, keep 48 PCA directions, cut each 48-d vector into 6 blocks of 8 numbers and replace each block by the nearest of 256 centroids learned by k-means on that block. A word is then 6 bytes (6 centroid ids), plus a small int8 codebook ( $6\times256\times8$ numbers). Plain scalar quantisation of fewer PCA dimensions at the same size was clearly worse (figure below): for a fixed budget, product quantisation keeps more dimensions, and dimensions matter more than precision.
Fast updates by branch and bound. Updating the penalties of all 10 001 words after every answer is too slow for 7 ms per round. Since penalties only ever increase, we keep an active set of the 300 best candidates and remember $LB$ , the smallest penalty outside it. Only active words are updated after each answer; when the best active penalty exceeds $LB$ , the outside words are brought up to date in one matrix product and the active set is rebuilt. The guess is still the exact argmin, at a fraction of the cost.
CPU budget. Each round gets a fair share of the remaining CPU time; past it, the program freezes the ranking and walks down it without further updates.

The core update, for each answer:

dv = X[g] - X[other]                    # cos(t,g) > cos(t,other)  <=>  t . dv > 0
new.append(dv if h == "hotter" else -dv)
pen[A] += np.logaddexp(0, -BETA * (X[A] @ np.array(new).T)).sum(1)   # active set only

Bar chart: mean score per round for several ways of storing the word vectors. Full 300-dimensional floats reach 92.9; among the formats that fit in the file, product quantisation of 48 PCA dimensions into 6 bytes gives the best score, 82.6, ahead of 20 PCA dimensions at 3 bits (80.8) and 16 dimensions at 4 bits (77.5).

Results

All numbers come from our simulated judge (Common Crawl vectors), never from Kattis. Rows 2 to 6 use 150 rounds whose target has a vector in both models and no time limit; the last rows are the generated program in a real interactive loop, on targets drawn from the whole vocabulary.

approach	median guesses	mean score per round
baseline: alphabetical order	5000	7.1 (exact expectation)
upper bound: our model equals the judge's	13	98.3
Wikipedia vectors, full 300-d, soft penalty	46	92.9
same, hard count of violated cuts	61	88.7
same, hinge penalty	58	89.6
PCA 16-d x 4 bits (8 bytes per word)	126	77.5
PQ 48-d, 6 bytes per word (chosen)	82	82.6
failed idea: hard pruning of unlikely candidates (for speed)	934	38.8
final program, group 1 setting ( $N=10$ , 5 runs)	82 to 158	77.5
final program, group 2 setting ( $N=100$ )	98	74.3
final program, group 3 time budget (CPU share of $N=1000$ , 200 rounds)	2511	20.7

Histogram of the number of guesses over 100 simulated rounds, on a log scale, with the per-round score curve. Most rounds need 30 to 300 guesses and score 70 to 95 points; a tail of rounds needing 1000 to 10000 guesses scores almost nothing.

Weighted over the three groups (30/40/30), this is about $0.3\times77.5 + 0.4\times74.3 + 0.3\times20.7 \approx 59$ points out of 100 against our simulated judge. Two honest warnings. First, the long rounds come mostly from words our vectors know badly (or not at all), and they also eat the CPU budget: with $N=1000$ even the output of thousands of guesses is costly in Python, so group 3 may run into the time limit on Kattis. Second, the real judge's vectors may be closer to ours (better scores) or farther (worse); the method itself does not depend on which.

The pruning row is a lesson in itself: throwing away candidates that look unlikely made each step faster but regularly threw away the target, because our model disagrees with the judge's. The branch-and-bound active set gives the same speed-up without changing a single decision.

What to remember

Turn feedback into geometry. "Closer to $g$ than to $b$ " is the half-space $\mathbf{t}\cdot(\mathbf{g}-\mathbf{b})>0$ ; any comparison answer can be read this way.
Use soft, probabilistic evidence when your model is imperfect. Summing log-likelihoods (a logistic noise model) degrades gracefully; hard elimination fails as soon as one cut is wrong.
When the model must ship inside the program, compression is part of the solution. PCA plus product quantisation keeps 82.6 of the 92.9 points at 1/200 of the size.
Exact speed-ups beat approximate ones. Penalties only grow, so a lower bound lets you update a small active set and still pick the true best candidate.
If you cannot see the judge, build one and say so. A simulated oracle with a different model than yours gives an honest estimate of the gap between "our model" and "their model".

"""SAIO "Hotter or Colder": build a self-contained Kattis submission.

The judge's similarity model is hidden and the program runs offline, so we ship
our OWN word vectors inside the source file (Kattis limit: 128 KiB in total).

  1. Our model: fastText Swedish Wikipedia vectors (wiki.sv.vec, first ~300 MB,
     filtered to the 10 001 vocabulary words; 533 words are missing and are
     guessed last).
  2. Compression: centre + PCA to 48 dims, then product quantisation (6 blocks of
     8 dims, 256 centroids each) -> 6 bytes per word, + int8 codebooks.
  3. Everything is lzma-compressed and base85-encoded into the generated program.

The generated program keeps a score (log-likelihood) for every candidate target t.
Each answer "g is hotter than b" says cos(t,g) > cos(t,b): a half-space constraint
on t. We add log sigmoid(beta * (cos(t,g) - cos(t,b))) (or its mirror for "colder")
and always guess the most likely unguessed word. For speed, only an active set of
the 300 best candidates is updated online (branch and bound: penalties only grow,
so outside words are refreshed only when they could become the best).

Usage:  python solution.py wiki.npy vocabulary.txt  ->  writes submission.py
(wiki.npy = (10001, 300) float array aligned with vocabulary.txt, zeros if missing)
"""
import sys, lzma, base64
import numpy as np
from sklearn.cluster import KMeans

D, M, K = 48, 6, 256            # PCA dims, PQ blocks, centroids per block

TEMPLATE = r'''import sys, os, lzma, time
from base64 import b85decode as d
import numpy as np
T0 = time.process_time()
def u(s): return lzma.decompress(d(s))
W = u(VOC).decode().split("\n"); V = len(W); idx = {w: i for i, w in enumerate(W)}
codes = np.frombuffer(u(COD), np.uint8).reshape(V, @M@)
books = np.frombuffer(u(BOO), np.int8).reshape(@M@, 256, -1).astype(np.float32) * @SC@
X = np.concatenate([books[j][codes[:, j]] for j in range(@M@)], 1)
ok = np.unpackbits(np.frombuffer(u(OK), np.uint8))[:V].astype(bool)
X[~ok] = 0; n = np.linalg.norm(X, axis=1); n[n == 0] = 1; X /= n[:, None]
BETA, K = np.float32(@BETA@), @K@
def sp(z): return np.logaddexp(0, z)                 # softplus = -log sigmoid(-z)
inp = sys.stdin.readline
N = int(inp())
LIMIT = 7.0                                          # CPU seconds we allow ourselves (limit 8)
NB = int(os.environ.get("BUDGET_N", N))              # (local tests only: pretend N is larger)
for rnd in range(N):
    start = inp().strip()
    t_round = time.process_time()
    budget = (LIMIT - (t_round - T0)) / max(1, NB - rnd)   # fair share of the remaining CPU
    pen = np.where(ok, 0.0, 1e6).astype(np.float32)  # -log likelihood of each candidate
    used = np.zeros(V, bool)
    prev = best = idx.get(start)
    if prev is not None: used[prev] = True
    dirs, c0, A, LB, order = [], 0, np.zeros(0, int), -1.0, None
    while True:
        if order is None and time.process_time() - t_round > budget:
            # out of time: freeze the ranking, then just walk down it
            if len(dirs) > c0:
                inA = np.zeros(V, bool); inA[A] = True
                st = np.where(~used & ~inA)[0]
                pen[st] += sp(-BETA * (X[st] @ np.array(dirs[c0:]).T)).sum(1)
            order = iter(np.argsort(pen))
        if order is not None:
            g = next(order)
            while used[g]: g = next(order)
        else:
            if len(A) == 0 or pen[A].min() > LB:
                # branch and bound: penalties only grow, so words outside the active
                # set A cannot beat LB; bring them up to date only when needed
                if len(dirs) > c0:
                    inA = np.zeros(V, bool); inA[A] = True
                    st = np.where(~used & ~inA)[0]
                    pen[st] += sp(-BETA * (X[st] @ np.array(dirs[c0:]).T)).sum(1)
                c0 = len(dirs); rest = np.where(~used)[0]
                if len(rest) > K:
                    part = np.argpartition(pen[rest], K); A = rest[part[:K]]; LB = pen[rest[part[K]]]
                else:
                    A, LB = rest, np.inf
            j = int(np.argmin(pen[A])); g = int(A[j]); A = np.delete(A, j)
        g = int(g); used[g] = True
        print(W[g], flush=True)
        r = inp().split()
        if len(r) < 2:                               # YES / NO: round over
            break
        new = []
        for other, h in ((prev, r[0]), (best, r[1])):
            if other is None: continue
            dv = X[g] - X[other]                     # cos(t,g) > cos(t,other)  <=>  t . dv > 0
            new.append(dv if h == "hotter" else -dv)
        if order is None and new:
            dirs += new
            if len(A): pen[A] += sp(-BETA * (X[A] @ np.array(new).T)).sum(1)
        if r[1] == "hotter" or best is None: best = g
        prev = g
print(round(time.process_time() - T0, 2), file=sys.stderr)
'''


def b85(b):
    return base64.b85encode(lzma.compress(b, preset=9 | lzma.PRESET_EXTREME)).decode()


def compress(E, ok):
    X = E / np.maximum(np.linalg.norm(E, axis=1, keepdims=True), 1e-9)
    mu = X[ok].mean(0)
    _, _, Vt = np.linalg.svd(X[ok] - mu, full_matrices=False)
    Z = (X - mu) @ Vt[:D].T
    s = D // M
    codes = np.zeros((len(Z), M), np.uint8); books = np.zeros((M, K, s), np.float32)
    for j in range(M):
        km = KMeans(K, n_init=1, random_state=0).fit(Z[ok][:, j*s:(j+1)*s])
        codes[:, j] = km.predict(Z[:, j*s:(j+1)*s]); books[j] = km.cluster_centers_
    scale = np.abs(books).max() / 127
    return codes, np.round(books / scale).astype(np.int8), scale


def fill_missing(E, vocab, minp=4):
    """A word without a vector gets the mean vector of the known words sharing its
    longest prefix (>= minp letters): Swedish inflections share their stem."""
    ok = np.linalg.norm(E, axis=1) > 0; E = E.copy()
    known = [(w, i) for i, w in enumerate(vocab) if ok[i]]
    for i, w in enumerate(vocab):
        if ok[i]: continue
        for L in range(len(w) - 1, minp - 1, -1):
            m = [j for v, j in known if v[:L] == w[:L]]
            if m:
                E[i] = np.mean([E[j] / np.linalg.norm(E[j]) for j in m], 0); break
    return E


def build(E, vocab, beta=20.0, k=300):
    E = fill_missing(E, vocab)
    ok = np.linalg.norm(E, axis=1) > 0
    codes, books, scale = compress(E, ok)
    src = (TEMPLATE.replace("@M@", str(M)).replace("@SC@", repr(float(scale)))
           .replace("@BETA@", str(beta)).replace("@K@", str(k)))
    head = (f'VOC = "{b85("\n".join(vocab).encode())}"\n'
            f'COD = "{b85(codes.tobytes())}"\n'
            f'BOO = "{b85(books.tobytes())}"\n'
            f'OK = "{b85(np.packbits(ok).tobytes())}"\n')
    return head + src


if __name__ == "__main__":
    E = np.load(sys.argv[1]).astype(np.float32)
    vocab = [w.strip() for w in open(sys.argv[2], encoding="utf-8") if w.strip()]
    src = build(E, vocab)
    open("submission.py", "w", encoding="utf-8").write(src)
    print("submission.py:", len(src.encode()), "bytes (Kattis limit 131072)")