My little tokenizer can't be this greedy?!

15s 1024 MB
Exclusive Grandmaster 10
AlgorithmsData Structures +1

Problem contributed by @root

Everyone congratulated the model for learning the ancient ciphered language. Nobody congratulated the tokenizer, who did all the work, deterministically, in one pass, without a single gradient.

Today you are the tokenizer.

Implement byte-pair encoding (BPE) training and encoding.

Training

The corpus is a sequence of n bytes. The initial token sequence is those bytes as token ids 0..255. Then repeat, up to k times:

  1. For every ordered pair (A, B), its count is the number of positions i in the current sequence with t[i] = A and t[i+1] = B. (Positions may overlap: the run aaa gives (a,a) a count of 2.)
  2. Let (A, B) be the pair with the largest count, breaking ties by smallest A, then smallest B. If the largest count is < 2, training stops.
  3. Record the merge. The new token's id is 256 + j where j is the 0-indexed merge number.
  4. Replace occurrences of (A, B) by the new token in a single left-to-right greedy scan: walk i from the start; if (t[i], t[i+1]) = (A, B), emit the new token and advance by 2, otherwise emit t[i] and advance by 1. (So aaaa under (a,a) → Z becomes ZZ, not aZa.)

Note the spec is declarative: after each merge, counts are simply recomputed on the resulting sequence. How you make that fast is your problem.

Encoding

Each query is a sequence of bytes, starting as tokens 0..255. Repeat: among all adjacent pairs currently present in the query that appear in the learned merge table, apply the one with the smallest merge index — replacing all its occurrences with the same left-to-right greedy scan — until no adjacent pair in the query is in the table. (These are exactly GPT-2 bpe() semantics.)

Byte blocks are hex-encoded (two lowercase hex digits per byte) so that tests are plain ASCII:

Text 
n k q
<2n hex chars>
len_1
<2*len_1 hex chars>
...
len_q
<2*len_q hex chars>

Bytes may take any value 00–ff. Input can be ~45 MB; use fast I/O.

A single u64: the FNV-1a hash of the full result, printed as unsigned decimal.

C 
uint64_t h = 0xcbf29ce484222325;           // fold(w): h = (h ^ w) * 0x100000001b3
fold(M);                                    // number of merges performed
for j in 0..M-1: fold(left_j); fold(right_j);
fold(final corpus length in tokens); for each corpus token t: fold(t);
fold(q); for each query: fold(its token count); for each token t: fold(t);
  • 1 <= n <= 2^23, 0 <= k <= 2^15
  • 0 <= q <= 16; each query at most 2^20 bytes, at most 2^22 query bytes total

Based on IOAI 2024's at-home NLP task (classification in an invented enciphered language).

Accepted 2/6
Acceptance 33%
Next up
The NOexistenceN of you AND me Similar problem
Insane 7
Escape the Matrix Similar problem
Normal+ 4
One step forward, two steps back Similar problem
Normal+ 4
Loading editor...
Input 
1 5 1
61
1
61
Expected Output 
7052655582604820982