代码

#include "router.h"

#include <stdint.h>
#include <netinet/in.h>

#include <new>
#include <unordered_set>
#include <algorithm>

inline uint32_t mask(unsigned char len) {
    return len == 0 ? 0 : ((~0) << (32 - len));
}

struct Node {
    uint32_t prefix;
    unsigned char len;
    bool valid;
    uint32_t nexthop;
    Node *next[2];
} *pool = new Node[2 << 20];

inline Node *alloc() {
    return pool ++;
}

Node *root = 0;

Node **trie_find(uint32_t addr, unsigned char lenlim = 32, bool insert = false) {
    Node **cur = &root, **val = 0;
    while (*cur) {
        if ((addr & mask((*cur)->len)) != (*cur)->prefix) break;
        if ((*cur)->valid) val = cur;
        if ((*cur)->len >= lenlim) break;

        size_t b = (addr & (1 << (31 - (*cur)->len))) ? 1 : 0;
        cur = &((*cur)->next[b]);
    }
    return insert ? cur : val;
}

void trie_add(uint32_t prefix, unsigned char len, uint32_t nexthop) {
    Node **f = trie_find(prefix, len, true);

    if (*f == 0) {
        *f = new(alloc()) Node {
            .prefix = prefix,
            .len = len,
            .valid = 1,
            .nexthop = nexthop,
            .next { 0, 0 }
        };
    } else {
        uint32_t pattern = (prefix ^ (*f)->prefix) | ~mask(std::min((*f)->len, len));
        unsigned char common = pattern == 0 ? 32 : __builtin_clz(pattern);

        size_t pb = (pattern & (1 << (31 - common))) ? 1 : 0;

        if (common == (*f)->len) {
            (*f)->valid = true;
            (*f)->nexthop = nexthop;
        } else {
            Node *new_node = new(alloc()) Node {
                .prefix = prefix & mask(common),
                .len = common,
                .valid = 0,
                .next { 0, 0 }
            };
            new_node->next[!pb] = *f;
            *f = new_node;
            if (common == len) {
                new_node->valid = true;
                new_node->nexthop = nexthop;
            } else {
                new_node->next[pb] = new(alloc()) Node {
                    .prefix = prefix,
                    .len = len,
                    .valid = 1,
                    .nexthop = nexthop,
                    .next { 0, 0 }
                };
            }
        }
    }
}

void init(int n, int q, const RoutingTableEntry *a) {
	for (size_t i = 0; i < (size_t) n; i ++) {
        trie_add(ntohl(a[i].addr), a[i].len, a[i].nexthop);
    }
}

unsigned query(unsigned addr) {
    Node **node = trie_find(ntohl(addr));
    if (! node || ! *node) return 0;
    return (*node)->nexthop;
}

评测结果