Paint Me a Flag

The solver starts one union-find element per pixel, with size = 1 and hash = mix(p ^ 0x243f6a8885a308d3) on each. It then makes two passes over the canvas. The first pass unions every pixel with its right and down neighbour when they share a colour, which collapses each same-colour region to a single root with the correct size and XOR hash. The second pass walks the same right and down edges again and, whenever two edges connect different roots, records them as neighbours in an adjacency set. After that any root with more than 48 neighbours is promoted to heavy. From there the simulation is just repeated fills, and because fills only ever recolour a root and merge it with same-colour neighbours, the partition only coarsens, which is exactly the one-directional behaviour union-find handles well.

Storing adjacency as a set that starts as a small vector and upgrades to a hash map once it grows past 32 entries keeps the common case (tiny regions with a few neighbours) cheap while still letting the background blob hold thousands of neighbours without quadratic scans.

My first version kept, for every component, a count of boundary edges per neighbour colour, which made each round’s “does a neighbour have colour c” check O(1). But every real operation recolours the background blob, and the blob borders thousands of components (about 3,700 on average early in part 2). Each recolour then had to update all of those neighbours’ counters. That version managed roughly 140k operations per minute and got slower as the blob grew.

The fix is a heavy/light split, a standard trick for graphs with a few very high-degree nodes:

• A component with more than 48 neighbours is heavy. Each heavy component keeps a lazy colour -> [neighbour roots] index, and stale entries (neighbours that have since merged or changed colour) are dropped the next time that list is read.
• When any component changes colour, it notifies only its heavy neighbours. There are few of those, and each component keeps a short list of which heavy components border it.
• A light component answers neighbour queries by scanning its own small adjacency set directly.

So when the blob recolours, it touches only a handful of heavy neighbours instead of thousands. Merging the blob with a light neighbour costs time proportional to the light side’s adjacency, which is the small-to-large merging rule that keeps the total work near-linear.

// Fast reimplementation of paint_me_a_flag (FortID rev challenge).
// usage: pmaf START.raw OPERATIONS.pmaf OUT.raw
// START/OUT are raw RGB24 dumps (W*H*3 bytes); PNG conversion is done in Python.
//
// Each op is a bucket fill. The op stream is XORed with a splitmix64 keystream
// whose state is chained on every fill's (component size, XOR-hash, old colour).
// The original BFS-fills 1 + 2*rounds times per op, over components that grow
// to most of the canvas. Here same-colour connected components live in a
// union-find (components only ever merge, never split), so a fill is just
// "recolour root, merge with neighbours of the new colour".
//
// Neighbour-colour queries use a heavy/light split: components with many
// neighbours ("heavy") keep a lazy colour -> neighbour-list index that their
// neighbours push into when they recolour; light components are scanned.
#include <cstdio>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <vector>
#include <unordered_map>
#include <utility>
using namespace std;
typedef uint64_t u64;
typedef uint32_t u32;

static const u64 B = 0xbf58476d1ce4e5b9ULL, P = 0x94d049bb133111ebULL;
static const u64 G = 0x9e3779b97f4a7c15ULL, C1 = 0xd6e8feb86659fd93ULL;
static const u64 HK = 0x243f6a8885a308d3ULL;
static inline u64 mix(u64 z) { z ^= z >> 30; z *= B; z ^= z >> 27; z *= P; z ^= z >> 31; return z; }

// set of u32: linear vector when small, hash map when large
struct USet {
    vector<u32> v;
    unordered_map<u32, char> *h = nullptr;
    size_t size() const { return h ? h->size() : v.size(); }
    bool has(u32 k) const {
        if (h) return h->count(k);
        for (u32 x : v) if (x == k) return true;
        return false;
    }
    void add(u32 k) {
        if (h) { (*h)[k] = 1; return; }
        if (has(k)) return;
        v.push_back(k);
        if (v.size() > 32) { h = new unordered_map<u32, char>(); h->reserve(128); for (u32 x : v) (*h)[x] = 1; vector<u32>().swap(v); }
    }
    void erase(u32 k) {
        if (h) { h->erase(k); return; }
        for (size_t i = 0; i < v.size(); i++) if (v[i] == k) { v[i] = v.back(); v.pop_back(); return; }
    }
    template <class F> void each(F f) const {
        if (h) { for (auto &p : *h) f(p.first); }
        else for (u32 x : v) f(x);
    }
    void clear() { delete h; h = nullptr; vector<u32>().swap(v); }
};

static const size_t HEAVY = 48;
static u32 W, H, N;
static vector<u32> par, sz, col, stamp;
static vector<u64> hs;
static vector<USet> adj;                                   // exact neighbour roots
static vector<unordered_map<u32, vector<u32>> *> cidx;      // heavy only: colour -> neighbour roots (lazy, may be stale)
static vector<vector<u32>> hnb;                             // heavy neighbours (lazy, may be stale)
static u32 curStamp = 0;

static u32 findr(u32 x) { while (par[x] != x) { par[x] = par[par[x]]; x = par[x]; } return x; }
static bool heavy(u32 r) { return cidx[r] != nullptr; }

static void makeHeavy(u32 r) {
    cidx[r] = new unordered_map<u32, vector<u32>>();
    adj[r].each([&](u32 m) { (*cidx[r])[col[m]].push_back(m); hnb[m].push_back(r); });
}

// distinct current roots of colour c adjacent to r (compacts r's lazy list)
static void nbrsOfColour(u32 r, u32 c, vector<u32> &out, bool firstOnly) {
    out.clear();
    ++curStamp;
    if (heavy(r)) {
        auto it = cidx[r]->find(c);
        if (it == cidx[r]->end()) return;
        vector<u32> &L = it->second;
        size_t k = 0;
        for (size_t i = 0; i < L.size(); i++) {
            u32 m = findr(L[i]);
            if (m == r || col[m] != c || stamp[m] == curStamp) continue;
            stamp[m] = curStamp;
            L[k++] = m;
            out.push_back(m);
        }
        L.resize(k);
        if (k == 0) cidx[r]->erase(it);
    } else {
        adj[r].each([&](u32 m) { if (col[m] == c && (!firstOnly || out.empty())) out.push_back(m); });
    }
}

// fold root b into root a (adjacent, same colour); returns new root
static u32 merge(u32 a, u32 b) {
    bool ha = heavy(a), hb = heavy(b);
    if ((hb && !ha) || (ha == hb && adj[a].size() < adj[b].size())) { swap(a, b); swap(ha, hb); }
    adj[a].erase(b);
    adj[b].erase(a);
    adj[b].each([&](u32 m) {
        adj[m].erase(b);
        adj[m].add(a);
        adj[a].add(m);
        if (ha && !hb) { (*cidx[a])[col[m]].push_back(m); hnb[m].push_back(a); }
    });
    if (hb) {
        for (auto &kv : *cidx[b]) { auto &d = (*cidx[a])[kv.first]; d.insert(d.end(), kv.second.begin(), kv.second.end()); }
        delete cidx[b]; cidx[b] = nullptr;
    }
    for (u32 h : hnb[b]) { u32 hr = findr(h); if (hr != a && hr != b) hnb[a].push_back(hr); }
    vector<u32>().swap(hnb[b]);
    adj[b].clear();
    par[b] = a; sz[a] += sz[b]; hs[a] ^= hs[b];
    if (hnb[a].size() > 4 * adj[a].size() + 64) {  // dedup lazy heavy-neighbour list
        ++curStamp; size_t k = 0;
        for (u32 h : hnb[a]) { u32 hr = findr(h); if (hr != a && heavy(hr) && stamp[hr] != curStamp) { stamp[hr] = curStamp; hnb[a][k++] = hr; } }
        hnb[a].resize(k);
    }
    if (!ha && adj[a].size() > HEAVY) makeHeavy(a);
    return a;
}

struct Res { u64 n, h; u32 c; };
static vector<u32> tmp;

static Res fill(u32 x, u32 nc) {
    u32 r = findr(x);
    Res res{sz[r], hs[r], col[r]};
    if (col[r] == nc) return res;
    col[r] = nc;
    ++curStamp;
    for (u32 h : hnb[r]) {  // tell heavy neighbours about the new colour
        u32 hr = findr(h);
        if (hr == r || !heavy(hr) || stamp[hr] == curStamp) continue;
        stamp[hr] = curStamp;
        (*cidx[hr])[nc].push_back(r);
    }
    nbrsOfColour(r, nc, tmp, false);
    vector<u32> same(tmp);
    for (u32 m : same) r = merge(r, m);
    return res;
}

static bool hasNbrColour(u32 r, u32 c) { nbrsOfColour(r, c, tmp, true); return !tmp.empty(); }

static u64 rd64(const unsigned char *p) { u64 v = 0; for (int i = 7; i >= 0; i--) v = (v << 8) | p[i]; return v; }
static u32 rd32(const unsigned char *p) { return p[0] | p[1] << 8 | p[2] << 16 | (u32)p[3] << 24; }

int main(int argc, char **argv) {
    if (argc != 4) { fprintf(stderr, "usage: %s start.raw ops.pmaf out.raw\n", argv[0]); return 2; }
    FILE *f = fopen(argv[2], "rb"); if (!f) { perror("ops"); return 1; }
    fseek(f, 0, SEEK_END); long fl = ftell(f); fseek(f, 0, SEEK_SET);
    vector<unsigned char> ops(fl); if (fread(ops.data(), 1, fl, f) != (size_t)fl) return 1; fclose(f);
    if (fl < 40 || memcmp(ops.data(), "PMAF\x03\r\n\x1a", 8)) { fprintf(stderr, "bad magic\n"); return 1; }
    W = rd32(&ops[8]); H = rd32(&ops[12]);
    u64 count = rd64(&ops[16]), seed = rd64(&ops[24]);
    u32 rounds = rd32(&ops[32]), resv = rd32(&ops[36]);
    N = W * H;
    fprintf(stderr, "W=%u H=%u count=%llu seed=%016llx rounds=%u resv=%u\n", W, H, (unsigned long long)count, (unsigned long long)seed, rounds, resv);
    if ((u64)fl != 40 + 8 * count) { fprintf(stderr, "size mismatch / trailing data\n"); return 1; }

    vector<unsigned char> img((size_t)N * 3);
    f = fopen(argv[1], "rb"); if (!f || fread(img.data(), 1, img.size(), f) != img.size()) { fprintf(stderr, "bad start raw\n"); return 1; } fclose(f);
    vector<u32> px(N);
    for (u32 i = 0; i < N; i++) px[i] = img[3*i] << 16 | img[3*i+1] << 8 | img[3*i+2];

    par.resize(N); sz.assign(N, 1); col = px; hs.resize(N); stamp.assign(N, 0);
    adj.resize(N); cidx.assign(N, nullptr); hnb.resize(N);
    for (u32 i = 0; i < N; i++) { par[i] = i; hs[i] = mix((u64)i ^ HK); }
    // union initial equal-coloured neighbours (plain DSU, adjacency built afterwards)
    auto uni = [&](u32 a, u32 b) { a = findr(a); b = findr(b); if (a == b) return; if (sz[a] < sz[b]) swap(a, b); par[b] = a; sz[a] += sz[b]; hs[a] ^= hs[b]; };
    for (u32 y = 0; y < H; y++) for (u32 x = 0; x < W; x++) {
        u32 p = y * W + x;
        if (x + 1 < W && px[p] == px[p+1]) uni(p, p+1);
        if (y + 1 < H && px[p] == px[p+W]) uni(p, p+W);
    }
    for (u32 y = 0; y < H; y++) for (u32 x = 0; x < W; x++) {
        u32 p = y * W + x;
        u32 nb[2] = {x + 1 < W ? p + 1 : UINT32_MAX, y + 1 < H ? p + W : UINT32_MAX};
        for (u32 q : nb) {
            if (q == UINT32_MAX) continue;
            u32 a = findr(p), b = findr(q);
            if (a != b) { adj[a].add(b); adj[b].add(a); }
        }
    }
    for (u32 i = 0; i < N; i++) if (par[i] == i && adj[i].size() > HEAVY) makeHeavy(i);

    u64 S = mix(mix(count) ^ seed ^ (u64)H ^ ((u64)W << 32) ^ (u64)rounds * C1 ^ 0x6f3d9b21a4c875e1ULL);
    u64 slow = 0;
    for (u64 i = 0; i < count; i++) {
        u64 ig = i * G;
        u64 a = ig ^ S;
        u64 d = mix(a) ^ rd64(&ops[40 + 8 * i]);
        u32 x = d & 0xffffff, y = (d >> 24) & 0xffffff;
        if ((u64)(d >> 48) != (mix((((u64)y << 24) | x) ^ a) >> 48)) { fprintf(stderr, "\nlost sync at op %llu\n", (unsigned long long)i); return 1; }
        if (x >= N) { fprintf(stderr, "\nbucket outside canvas at op %llu\n", (unsigned long long)i); return 1; }
        Res r0 = fill(x, y);
        S = mix(r0.n * C1 ^ mix(r0.h) ^ S ^ ig ^ mix(((u64)r0.c << 24) | y));
        for (u32 j = 1; j <= rounds; j++) {
            u64 sd = ig ^ S ^ (u64)j * C1;
            u32 c = (mix(sd) & 0xfffffe) | 1;
            while (c == 0xf2eadf) { sd = mix(sd + G); c = (mix(sd) & 0xfffffe) | 1; }
            u32 rt = findr(x);
            Res r1, r2;
            if (!hasNbrColour(rt, c)) {
                // fill C with c, then back with y: component unchanged, no merges
                r1 = {sz[rt], hs[rt], col[rt]};
                r2 = {sz[rt], hs[rt], c};
            } else {
                slow++;
                r1 = fill(x, c);
                r2 = fill(x, y);
            }
            u64 A = mix(r1.n * C1 ^ mix(r1.h) ^ ig ^ S ^ mix(((u64)r1.c << 24) | c));
            S = mix(r2.n * C1 ^ mix(r2.h) ^ ig ^ mix(((u64)r2.c << 24) | y) ^ A);
        }
        if ((i & 0xffff) == 0) fprintf(stderr, "\rop %llu/%llu slow=%llu", (unsigned long long)i, (unsigned long long)count, (unsigned long long)slow);
    }
    fprintf(stderr, "\ndone, slow-path rounds=%llu\n", (unsigned long long)slow);
    for (u32 i = 0; i < N; i++) { u32 c = col[findr(i)]; img[3*i] = c >> 16; img[3*i+1] = c >> 8; img[3*i+2] = c; }
    f = fopen(argv[3], "wb"); fwrite(img.data(), 1, img.size(), f); fclose(f);
    return 0;
}

"""Convert PNG<->raw RGB and drive the fast pmaf solver for each flag part.

usage: python -I run.py HANDOUT_DIR OUT_DIR SOLVER_BIN [part numbers...]
"""
import subprocess
import sys
from pathlib import Path

from PIL import Image

handout, out, solver = Path(sys.argv[1]), Path(sys.argv[2]), sys.argv[3]
parts = [int(p) for p in sys.argv[4:]] or range(1, 14)
out.mkdir(parents=True, exist_ok=True)

for n in parts:
    start = Image.open(handout / f"flag_part_{n:02d}.png").convert("RGB")
    raw_in, raw_out = out / f"start_{n:02d}.raw", out / f"out_{n:02d}.raw"
    raw_in.write_bytes(start.tobytes())
    subprocess.run([solver, str(raw_in), str(handout / f"flag_part_{n:02d}.pmaf"), str(raw_out)], check=True)
    Image.frombytes("RGB", start.size, raw_out.read_bytes()).save(out / f"out_{n:02d}.png")
    raw_in.unlink()
    raw_out.unlink()
    print(f"part {n:02d} -> {out / f'out_{n:02d}.png'}", flush=True)