代码

#include "router.h"

#include <stdint.h>
#include <arpa/inet.h>

const int MAX_N = 900000;
const int MAX_N_NODES = MAX_N * 32;

struct TrieNode {
	TrieNode *s[2];
	uint32_t nexthop;
	uint32_t prefix;
	int prefix_len;
};

static TrieNode nodes[MAX_N_NODES], *last_node = nodes, *root;

static inline TrieNode * new_node() {
	return last_node++;
}

static void trie_insert(TrieNode *root, uint32_t addr, int len, uint32_t nexthop) {
	for (int i = 0; i < len; i++) {
		if (i == root->prefix_len) {
			TrieNode *&t = root->s[((addr >> (31 - i))) & 1];
			if (!t) {
				t = new_node();
				t->prefix = addr;
				t->prefix_len = len;
			}
			root = t;
		} else if ((((root->prefix ^ addr) >> (31 - i)) & 1) != 0) {
			int temp = ((addr >> (31 - i))) & 1;
			TrieNode *t1 = new_node(), *t2 = new_node();
			*t1 = *root;
			root->s[temp ^ 1] = t1;
			root->s[temp] = t2;
			root->nexthop = 0;
			root->prefix_len = i;
			t2->prefix = addr;
			t2->prefix_len = len;
			root = t2;
		}
	}
	
	root->nexthop = nexthop;
}

static uint32_t trie_query(TrieNode *root, uint32_t addr) {
	uint32_t ans = 0;
	
	for (int i = 0; i < 32; i++) {
		if (i == root->prefix_len) {
			if (root->nexthop) ans = root->nexthop;
			TrieNode *t = root->s[((addr >> (31 - i))) & 1];
			if (!t) return ans;
			root = t;
		} else if ((((root->prefix ^ addr) >> (31 - i)) & 1) != 0) {
			return ans;
		}
	}
	
	return root->nexthop ? root->nexthop : ans;
}

void init(int n, int q, const RoutingTableEntry *a) {
	root = new_node();
	for (int i = 0; i < n; i++) {
		trie_insert(root, htonl(a[i].addr), a[i].len, a[i].nexthop);
	}
}

unsigned query(unsigned addr) {
	return trie_query(root, htonl(addr));
}

评测结果