BPE Pair Counts & Merge

The problem

BPE (byte-pair encoding) builds a vocabulary by repeatedly gluing the most frequent pair of neighbouring tokens into a new token (see the theory page). Each round needs two operations:

count how often every pair of neighbours occurs;
merge: replace every occurrence of the chosen pair with a single new id.

You will write both. They work on a plain Python list of integer token ids. The next two exercises are built on them.

Part 1: counting pairs

Walk through the list and look at every pair of adjacent ids: positions 0 and 1, then 1 and 2, and so on. A list of $n$ ids has $n - 1$ such pairs.

Example. [1, 2, 3, 1, 2] has the pairs (1, 2), (2, 3), (3, 1), (1, 2). So the counts are {(1, 2): 2, (2, 3): 1, (3, 1): 1}.

The dict must list the pairs in the order of their first occurrence: here (1, 2) first, then (2, 3), then (3, 1). A Python dict keeps the order in which keys are inserted, so this comes for free if you insert as you walk. The next exercise uses that order to break ties.

Part 2: merging a pair

Build a new list, scanning left to right. At each position, if the id there and the next one form the pair, write the new id and jump two positions ahead; otherwise copy the id and move one position.

Example. merge([1, 2, 3, 1, 2], (1, 2), 99):

positions 0, 1 are 1, 2: the pair. Write 99, jump to position 2;

position 2 is 3, and 3, 1 is not the pair: copy 3, move to position 3;

positions 3, 4 are 1, 2: write 99, jump to position 5, the end.

The result is [99, 3, 99].

Overlaps. In [5, 5, 5] the pair (5, 5) appears at positions 0–1 and 1–2, but the middle 5 cannot belong to both. The jump of two positions settles it: merge([5, 5, 5], (5, 5), 7) replaces positions 0–1, jumps to position 2, copies the last 5, and returns [7, 5].

What to implement

get_pair_counts(ids): a dict mapping every pair of adjacent ids (a, b) to how many times it occurs, in the order of first occurrence.
merge(ids, pair, new_id): a new list where every occurrence of pair is replaced by new_id, scanning left to right.

Requirements

Overlaps as in the example: after a replacement, continue after the pair.
merge does not modify ids.
Lists of length 0 or 1 have no pairs: the counts are {}, and merge returns a copy.

Hints

zip(ids, ids[1:]) walks every adjacent pair.
For merge, use a while loop with an index i that you advance by 2 after a match and by 1 otherwise. Check i + 1 < len(ids) before reading ids[i + 1].