IOAI Training: Learn AI by solving problems

The problem in brief

You receive an English text (letters, digits, a few punctuation marks, spaces and newlines, at most 500 000 characters) in which some characters were replaced by #. Print the text back with every # replaced by the character that was there, with the right case. Anything can be hidden, including spaces and punctuation, but never a newline. There are three test files: in one, at most one letter per word is hidden and spaces and punctuation are never touched; in the others, 10 % and 30 % of all characters are hidden at random. The statement is in Swedish; the sample is the opening of Dickens' A Christmas Carol. There is no training data and the time limit is 1 second.

How you are scored

Each file is worth $P = 60$ points. If the file contains $N$ hidden characters and you restore $M$ of them correctly, you get $\lfloor P\cdot M/N \rfloor$ points, so the total is out of 180. Changing a character that was not a # gives zero for that file, and newlines in your output are ignored. There is no threshold: every restored character counts.

The official test files are hidden, so we built our own: two public-domain novels from Project Gutenberg, A Tale of Two Cities and Pride and Prejudice, cut to 500 000 characters, restricted to the allowed alphabet, and masked in the three ways described above (in the first mode, half of the words get one hidden letter).

First look at the data

The sample shows the three difficulties:

a hidden letter inside a frequent word (The#e is There), easy for a human;
a hidden space (about#that) or punctuation mark (done fo## is done for.): the model must also decide where words end;
runs of several # (whate##r, Ma##ey), where the candidates multiply.

The most important observation is not in the sample but in the sizes: an input has up to 500 000 characters and only 10 to 30 % are hidden. The input is its own training set. A novel repeats its names (Marley appears unmasked several times), its words and its turns of phrase; learning from the text itself beats any general model of English that would have to fit in the source file.

A first baseline

Replace every # by the most frequent character. In the random modes that is the space, which restores 16 to 17 % of the hidden characters (one character in six is a space); in the one-letter-per-word mode, where spaces are never hidden, the best constant is e, which restores 11.1 %. Any real model has to use the context.

The key idea

Treat the problem as decoding with a language model. A character $n$ -gram model gives the probability of a character given the $k$ previous ones,

P(x_i \mid x_{i-k}\dots x_{i-1}) \approx \frac{\#(x_{i-k}\dots x_{i-1}\,x_i)}{\#(x_{i-k}\dots x_{i-1})},

with counts taken from all windows of the input that contain no #. When a long context was never seen, we back off to a shorter one, multiplying by 0.4 for each step ("stupid backoff"); this is crude but fast and works well with this much data.

The essential point is to use both sides of the gap. For a single hidden character $c$ at position $i$ , the best choice maximises

\log P(c \mid \text{left}) + \sum_{q=i+1}^{i+k} \log P(x_q \mid x_{q-k}\dots x_{q-1}),

because every $n$ -gram that contains position $i$ carries evidence: in S#ORY, a left-only model hesitates between St, Sh and a space, but the following characters ORY settle it. For a run of $L$ hidden characters the candidates grow like $75^L$ , so we keep the 8 best partial fillings at each position (beam search), then rescore the finished ones with the right context.

A second, word-level pass fixes what the character model gets wrong inside words. When a token's boundaries are certain (no # just before or after it), its pattern (Gor#on's) is matched against the intact words of the text, and among the matches we pick the one most often seen after the previous word (bigram count, smoothed with the word's own frequency). This brings in long-range knowledge, such as names that appear elsewhere in the book, that a 5-character window cannot see.

The full solution

Count every $n$ -gram of length 1 to 6 that contains no # (that is, contexts of up to 5 characters).
Scan the text; for each maximal run of #, run the beam search over the alphabet of characters seen in the text, scoring each filling by the forward model plus the log-probability of the next 5 known characters.
Word pass: tokenise with [A-Za-z0-9'#]+, and for each token that contained a # with certain boundaries, replace it by the best-matching intact word of the text if one exists.
Print the result; characters that were not # are never changed, so we can never lose a file by accident.

for sc, f in B:                       # finished candidate fillings of the run
    tot, h = sc, left + f
    for q in range(j, min(n, j + right)):
        if s[q] == '#': break         # stop at the next unknown
        tot += lp(h[-N:], s[q]); h += s[q]

Grouped bar chart of the share of hidden characters restored on our own test files. Character model only on A Tale of Two Cities: 81, 79 and 60 percent for the three masking modes; with the word pass: 86, 82 and 62; with the word pass on Pride and Prejudice: 89, 86 and 66.

Speed warning. Our Python implementation needs 8 to 50 seconds of CPU on a 500 000-character file (the 30 % file is the slowest: three times more gaps and longer runs). The algorithm is linear in the text length, but the Kattis limit is 1 second, so a real submission must be the same algorithm in C++ (hash maps of $n$ -gram counts, a fixed-size beam). The accuracy numbers below are those of the algorithm, not a claim that the Python file passes the time limit.

Results

Share of hidden characters restored, and points out of 60 per file, on our self-built tests:

approach	one letter per word	10 % hidden	30 % hidden	total /180
constant character (e, then space)	11.1 % (6)	16.5 % (9)	16.7 % (10)	25
char 5-gram, left and right context (Two Cities)	81.2 % (48)	78.8 % (47)	59.6 % (35)	130
+ word pass (Two Cities)	86.5 % (51)	82.2 % (49)	61.5 % (36)	136
+ word pass (Pride and Prejudice)	89.2 % (53)	85.8 % (51)	66.3 % (39)	143
failed idea: also use the bigram with the next word (Two Cities)	85.2 % (51)
longer contexts, $N = 8$ (char only, Two Cities)	79.0 % (47)

Two lessons from the failed rows. Longer contexts are worse: with a fixed amount of text, 8-character contexts are seen too rarely and the crude backoff takes over. And adding the next word to the word pass hurt slightly, because when 30 to 50 % of neighbouring words are themselves damaged the extra evidence is often unreliable and our ad hoc combination over-weights it. The remaining errors are mostly genuinely ambiguous: #he (the or she), #e (he, we, me), capital letters in rare uppercase headings, and long runs at 30 % masking where whole words disappear.

What to remember

The input can be the training set. When the test document is large and mostly intact, statistics learned from it beat any general model you could fit in the program.
Use the context on both sides of a gap: every $n$ -gram that overlaps the missing character carries evidence, and the right side is as informative as the left.
Beam search keeps exponential searches affordable: keep a few best partial solutions instead of all $75^L$ fillings.
Combine levels: a character model handles spaces and punctuation, a word model brings long-range memory (names, rare words).
Check the time limit early. An algorithm that is right but 20 times too slow in Python needs a compiled implementation, not a cleverer model.

"""SAIO "Ratta fel": restore the characters hidden behind '#' in an English text.

No training data is given, and none is needed: the input itself (up to 500 000
characters, 70 to 95 % of them intact) is a large sample of the very text we
must repair. Two passes:

  1. Character n-gram model learned on the #-free windows of the input
     (contexts of up to N = 5 characters, "stupid backoff" with factor 0.4).
     Each run of '#' is filled by a small beam search (width 8); candidate
     fillings are rescored with the log-probability of the N characters that
     FOLLOW the run, so both the left and the right context are used.
  2. Word pass: inside a token whose boundaries are certain, the pattern
     (e.g. "Gor#on's") is matched against the intact words of the text and the
     best match by bigram count with the previous word wins.

Usage: python solution.py < input.txt > output.txt
       python solution.py --eval masked.txt original.txt   (accuracy on # only)
Note: this Python version takes 8 to 50 s on 500 000 characters; to pass the
1 s Kattis limit the same algorithm must be ported to C++ (counts in hash maps).
"""
import sys, math, collections, time, re
def train(t, N):
    """counts[k][ctx+ch] for contexts of length k, from #-free windows of the input itself."""
    cnt=[collections.Counter() for _ in range(N+1)]; ctx=[collections.Counter() for _ in range(N+1)]
    for k in range(N+1):
        c,cx=cnt[k],ctx[k]
        for i in range(k,len(t)):
            w=t[i-k:i+1]
            if '#' in w: continue
            c[w]+=1; cx[w[:-1]]+=1
    return cnt,ctx
def make_lp(cnt,ctx,N,V):
    cache={}
    def lp(h,ch):
        key=(h,ch)
        r=cache.get(key)
        if r is not None: return r
        pen=0.0
        for k in range(min(N,len(h)),-1,-1):
            hh=h[len(h)-k:] if k else ''
            if '#' in hh: continue
            c=cnt[k].get(hh+ch,0)
            if c: r=pen+math.log(c/ctx[k][hh]); break
            if ctx[k].get(hh,0): pen+=math.log(0.4)
        else: r=pen+math.log(1.0/V)
        cache[key]=r; return r
    return lp
def solve(t, N=5, beam=8, right=None):
    right = N if right is None else right
    t0=time.time(); cnt,ctx=train(t,N); alpha=sorted(set(t)-{'#','\n'}); V=len(alpha); lp=make_lp(cnt,ctx,N,V)
    s=list(t); i=0; n=len(s)
    while i<n:
        if s[i]!='#': i+=1; continue
        j=i
        while j<n and s[j]=='#': j+=1          # run s[i:j]
        left=''.join(s[max(0,i-N):i])
        B=[(0.0,'')]
        for p in range(i,j):
            nb=[]
            for sc,f in B:
                h=(left+f)[-N:]
                for ch in alpha: nb.append((sc+lp(h,ch),f+ch))
            nb.sort(reverse=True); B=nb[:beam]
        best=None
        for sc,f in B:                          # add the evidence of the following chars
            tot=sc; h=(left+f)
            for q in range(j,min(n,j+right)):
                ch=s[q]
                if ch=='#': break
                if ch=='\n': ch=' '
                tot+=lp(h[-N:],ch); h+=ch
            if best is None or tot>best[0]: best=(tot,f)
        s[i:j]=list(best[1]); i=j
    return ''.join(s)


TOK=re.compile(r"[A-Za-z0-9'#]+")
def word_pass(orig, filled):
    """Second pass: inside words, prefer real words of the text (unigram + bigram counts)."""
    toks=[(m.start(),m.end()) for m in TOK.finditer(orig)]
    words=[orig[a:b] for a,b in toks]
    uni=collections.Counter(w for w in words if '#' not in w)
    bi=collections.Counter((p.lower(),w.lower()) for p,w in zip(words,words[1:]) if '#' not in p+w)
    # index vocabulary by (length, pattern with one wildcard position) for fast lookup
    bylen=collections.defaultdict(list)
    for w,c in uni.items(): bylen[len(w)].append((w,c))
    out=list(filled); cache={}
    for k,(a,b) in enumerate(toks):
        w=orig[a:b]
        if '#' not in w: continue
        # the token boundaries are certain only if the chars around are not '#'
        if (a>0 and orig[a-1]=='#') or (b<len(orig) and orig[b]=='#'): continue
        pat='^'+re.escape(w).replace('\\#','.').replace('#','.')+'$'
        key=pat
        if key not in cache:
            rx=re.compile(pat); cache[key]=[(v,c) for v,c in bylen[len(w)] if rx.match(v)]
        cands=cache[key]
        if not cands: continue
        prev=words[k-1].lower() if k else ''
        cur=''.join(filled[a:b])
        def sc(v,c): return math.log(bi.get((prev,v.lower()),0)+0.05*c+1e-9) + (0.5 if v==cur else 0)
        best=max(cands,key=lambda vc: sc(*vc))[0]
        out[a:b]=list(best)
    return ''.join(out)

def accuracy(masked, out, ref):
    a, b, c = (x.replace("\n", "") for x in (masked, out, ref))
    pos = [i for i, ch in enumerate(a) if ch == "#"]
    return sum(b[i] == c[i] for i in pos) / len(pos)


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--eval":
        t = open(sys.argv[2]).read(); ref = open(sys.argv[3]).read()
        t0 = time.process_time(); out = word_pass(t, solve(t))
        acc = accuracy(t, out, ref)
        print(f"fixed {100*acc:.2f} % of the #, {int(60*acc)}/60 points, cpu {time.process_time()-t0:.1f} s")
    else:
        t = sys.stdin.read()
        sys.stdout.write(word_pass(t, solve(t)))

Rätta fel