vote.c

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <limits.h>

typedef struct region {
    unsigned *data;
    unsigned width;
    unsigned height;
    unsigned pitch;
} region;

enum direction { HORIZONAL, VERTICAL };

inline unsigned *
value_at(region r, unsigned x, unsigned y)
{
    return &r.data[y * (r.width + r.pitch) + x];
}

unsigned
population(region r)
{
    unsigned total = 0;
    for (unsigned y = 0; y < r.height; y++)
        for (unsigned x = 0; x < r.width; x++)
            total += *value_at(r, x, y);
    return total;
}

void
split(region *a, region *b, enum direction dir, int v)
{
    region r = *a;
    switch (dir) {
        case HORIZONAL:
            a->height = v;
            b->height = r.height - v;
            a->width = b->width = r.width;
            a->pitch = b->pitch = r.pitch;
            b->data = value_at(r, 0, v);
            break;
        case VERTICAL:
            a->width = v;
            b->width = r.width - v;
            a->height = b->height = r.height;
            a->pitch = r.pitch + b->width;
            b->pitch = r.pitch + a->width;
            b->data = value_at(r, v, 0);
            break;
    }
}

bool
in_region(region r, unsigned *p)
{
    if (p >= r.data) {
        unsigned diff = p - (r.data - r.pitch);
        unsigned y = diff / (r.width + r.pitch);
        return y < r.height && (diff % (r.width + r.pitch)) >= r.pitch;
    } else
        return false;
}

void
colormap(unsigned nregion, region *regions, int colors[nregion])
{
    for (unsigned i = 0; i < nregion; i++)
        colors[i] = -1;
    unsigned alist[nregion][nregion];
    for (unsigned i = 0; i < nregion; i++)
        for (unsigned j = 0; j < nregion; j++)
            alist[i][j] = 0;
    // Build adjacency list (graph)
    for (unsigned i = 0; i < nregion; i++) {
        for (unsigned x = 0; x < regions[i].width; x++) {
            unsigned *below = value_at(regions[i], x, -1);
            unsigned *above = value_at(regions[i], x, regions[i].height);
            for (unsigned j = i + 1; j < nregion; j++) {
                if (in_region(regions[j], below) ||
                    in_region(regions[j], above))
                    alist[i][j] = alist[j][i] = 1;
            }
        }
        for (unsigned y = 0; y < regions[i].height; y++) {
            unsigned *below = value_at(regions[i], -1, y);
            unsigned *above = value_at(regions[i], regions[i].width, y);
            for (unsigned j = i + 1; j < nregion; j++) {
                if (in_region(regions[j], below) ||
                    in_region(regions[j], above))
                    alist[i][j] = alist[j][i] = 1;
            }
        }
    }
    // Find 4-coloring by brute force
    unsigned i = 0;
    while (i < nregion) {
        if (++colors[i] == 4) {
            colors[i] = -1;
            i--;
            continue;
        }
        bool valid = true;
        for (unsigned j = 0; valid && j < nregion; j++)
            if (alist[i][j] && colors[i] == colors[j])
                valid = false;
        if (valid)
            i++;
    }
}

void
print(region whole, unsigned nsub, region *sub)
{
    int colors[nsub];
    colormap(nsub, sub, colors);
    for (unsigned y = 0; y < whole.height; y++) {
        for (unsigned x = 0; x < whole.width; x++) {
            unsigned *p = value_at(whole, x, y);
            unsigned color = 0;
            for (unsigned i = 0; i < nsub; i++)
                if (in_region(sub[i], p))
                    color = i;
            printf("\x1b[%dm%2u\x1b[0m ",
                   (int[]){91, 92, 93, 94}[colors[color]],
                   *p);
        }
        putchar('\n');
    }
}

void
evenly_split(region *a, region *b)
{
    region r = *a;
    enum direction best_dir = 0;
    unsigned best_v = 0;
    unsigned best_diff = UINT_MAX;
    unsigned best_length = UINT_MAX;
    for (enum direction dir = 0; dir <= 1; dir++) {
        unsigned max = dir == HORIZONAL ? r.height : r.width;
        unsigned length = dir == HORIZONAL ? r.width : r.height;
        for (unsigned v = 0; v < max; v++) {
            region ra = r;
            region rb;
            split(&ra, &rb, dir, v);
            int pop_a = population(ra);
            int pop_b = population(rb);
            unsigned diff = labs(pop_a - pop_b);
            if (diff <= best_diff && length <= best_length) {
                best_dir = dir;
                best_v = v;
                best_diff = diff;
                best_length = length;
            }
        }
    }
    split(a, b, best_dir, best_v);
}

int
main(void)
{
    // Read input
    unsigned nlines, nrows, ncols;
    scanf("%u %u %u", &nlines, &nrows, &ncols);
    unsigned data[ncols * nrows];
    for (unsigned i = 0; i < nrows * ncols; i++)
        scanf("%u", &data[i]);

    // Construct regions
    struct region whole = {
        .data = data,
        .width = nrows,
        .height = ncols,
        .pitch = 0
    };
    struct region regions[nlines + 1];
    regions[0] = whole;
    unsigned nregions = 1;
    for (unsigned i = 0; i < nlines; i++) {
        unsigned next = 0;
        unsigned next_pop = population(regions[0]);
        // Find currently most populated region
        for (unsigned r = 1; r < nregions; r++) {
            unsigned pop = population(regions[r]);
            if (pop > next_pop) {
                next = r;
                next_pop = pop;
            }
        }
        evenly_split(&regions[next], &regions[nregions++]);
    }

    // Print results
    for (unsigned i = 0; i < nregions; i++)
        printf("%u ", population(regions[i]));
    putchar('\n');
    print(whole, nregions, regions);
    return 0;
}