[C++] [算法] [Linux] radix tree

今天看Linux內核代碼時，看到了radix tree，從書上大概地瞭解了radix tree的一些基本知識，包括它的結構和操作。由於Linux內核代碼不容易看懂，因此，打算自己家先實現一下radix tree，然後再看內核中的代碼，由於懶得去查radix tree的標準表示方式，所以，這裏radix tree的結構與Linux內核使用的類似。

首先，簡單說一下radix tree的結構。

radix tree的葉節點是存放在樹中的數據，內節點包含子節點的指針，每個內節點的孩子個數固定，Linux內核一般設置爲64，也就是內節點包含64個指針，內節點還包括本節點的有效孩子個數以及所在的層次，這裏爲了簡單，並沒有包含Linux內核內節點那麼多數據。

radix tree的使用主要是提高查詢效率，因此，這裏說一下radix tree的查詢操作：

每個內節點包含的指針個數是pow(2, n)，好處是可以通過提取索引值的n位來得到該層的索引。由於這裏採用的是pow(2, 6)，因此，每層的索引值是6位。比如下圖中的索引值131，轉換成16進制就是0x00000083，內節點中存放的有當前節點所在的層次（從下往上計算），那麼，提取第二層的索引值是2，因此取第二層的slot[2]轉到第一層，然後提取第一層的索引值，是3，因此，提取第一層的slot[3]，由於這是第一層了，於是，遍歷結束，所得到的slot[3]就是希望查詢的數據的指針。

《深入理解Linux內核》上關於radix tree的圖示：

看了圖應該很好理解了，直接貼代碼吧，由於時間有限，沒有對代碼進行嚴格的測試，只測試了插入和查找操作，因此，代碼有可能會有問題：

#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
using namespace std;

typedef unsigned int u_int;

//常量定義
const u_int RADIX_TREE_SLOTS = 64; //一個內節點包含的最多的孩子個數
const u_int RADIX_MAX_LEVEL = 5; //最大的層數
const u_int radix_level[RADIX_MAX_LEVEL + 1] = { 0x00000000,
	                                             0x0000003F,
	                                             0x00000FC0,
	                                             0x0003F000,
	                                             0x00FC0000,
	                                             0x3F000000 }; //層次掩碼，用於提取某一層上的索引值

//函數的返回值狀態
enum radix_status {
	INSERT_SUCC,
	INSERT_FAIL,
	DELETE_SUCC,
	DELETE_FAIL
};

//radix tree的內節點
struct __radix_tree_node {
	u_int height; //內節點的高度，從下往上算起
	u_int count; //有效子節點的個數
	void *slot[RADIX_TREE_SLOTS];
};

//radix tree的根節點
struct __radix_tree_root {
	u_int height; //整個樹的高度，不包括樹根
	__radix_tree_node *rnode;
};

//radix tree的頁節點，存儲在樹中的數據，這裏暫時將數據成員設置爲string
struct radix_tree_data {
	string str;
};

class radix_tree {
public:
	radix_tree();
	void radix_tree_stretch(); //將樹向下擴展一層
	void radix_tree_stretch_n(u_int); //將樹擴展到n層
	radix_status radix_tree_insert(u_int, radix_tree_data&); //插入一個索引值爲index的數據
	radix_status radix_tree_delete(u_int, radix_tree_data *&); //刪除索引值爲index的數據
	radix_tree_data* radix_tree_lookup(u_int); //查詢索引值爲index的數據
	void radix_tree_print(); //打印整棵樹
	~radix_tree();
private:
	u_int get_want_level(u_int index) //從給定的索引值獲取樹所期望的層次
	{
		u_int level = RADIX_MAX_LEVEL;
		while(level >= 0) {
			if(index & radix_level[level]) {
				return level;
			}
			--level;
		}
		return level;
	}

	u_int get_level_index(u_int index, u_int level) //從給定的索引值和層次，得到該層的索引值
	{
		return (index & radix_level[level]) >> ((level - 1) * 6);
	}

	__radix_tree_node* radix_tree_alloc_node() //分配一個樹的內節點
	{
		__radix_tree_node *pr_node = new __radix_tree_node;
		pr_node->height = 0;
		pr_node->count = 0;
		//for(u_int i = 0; i < RADIX_TREE_SLOTS; ++i)
			//pr_node->slot[i] = NULL;
		memset(pr_node->slot, 0, sizeof(pr_node->slot));
		return pr_node;
	}

	__radix_tree_node* radix_tree_get_node(u_int); //根據索引值得到該索引值所在的節點
	void __radix_tree_destroy_node(__radix_tree_node *&pr_node); 
	void radix_tree_print_node(__radix_tree_node *pr_node);
	radix_status __radix_tree_insert(u_int, u_int, radix_tree_data&);

	__radix_tree_root r_tree;
};

//構造函數，用於初始化radix tree
radix_tree::radix_tree()
{
	r_tree.height = 0;
	r_tree.rnode = NULL;
}

//銷燬某個內節點，並遞歸銷燬它的子節點
void radix_tree::__radix_tree_destroy_node(__radix_tree_node *&pr_node)
{
	if(pr_node == NULL)
		return;

	if(pr_node->height == 1) {
		delete pr_node;
		pr_node = NULL;
		return;
	}

	for(u_int i = 0; i < RADIX_TREE_SLOTS; ++i) {
			__radix_tree_destroy_node((__radix_tree_node*&)(pr_node->slot[i]));
	}

	delete pr_node;
	pr_node = NULL;
}

//析構函數，用於銷燬radix tree
radix_tree::~radix_tree()
{
	__radix_tree_destroy_node(r_tree.rnode);
}

//向下擴展一層：分配一個內節點，然後將這個新節點的第一個slot的指針指向根節點指向的節點
//最後更新新節點中的信息和整個樹的高度
void radix_tree::radix_tree_stretch()
{
    __radix_tree_node *pr_node = radix_tree_alloc_node();
    pr_node->height = r_tree.height + 1;
    pr_node->count = 1;
    pr_node->slot[0] = r_tree.rnode;
    r_tree.rnode = pr_node;
    ++r_tree.height;
    return;
}

//向下擴展到level層
void radix_tree::radix_tree_stretch_n(u_int level)
{
	u_int height = r_tree.height;
	if(height >= level) {
		return;
	}
	while(height < level) {
		radix_tree_stretch();
		++height;
	}
}

//插入索引值的輔助函數，採用此函數時，說明該樹此時的高度能夠插入index值的數據
radix_status radix_tree::__radix_tree_insert(u_int index, u_int max_level, radix_tree_data& data)
{
	u_int cur_index = 0;
	__radix_tree_node *pr_node = r_tree.rnode, *pr_node_tmp = NULL;
	while(max_level > 1) {
		cur_index = get_level_index(index, max_level);
		pr_node_tmp = (__radix_tree_node*)(pr_node->slot[cur_index]);
		if(pr_node_tmp == NULL) {
			pr_node->slot[cur_index] = radix_tree_alloc_node();
			((__radix_tree_node*)(pr_node->slot[cur_index]))->height = pr_node->height - 1;
			pr_node_tmp = (__radix_tree_node*)(pr_node->slot[cur_index]);
		}
		pr_node = pr_node_tmp;
		--max_level;
	}
	cur_index = get_level_index(index, 1);
	if(pr_node->slot[cur_index]) {
		return INSERT_FAIL;
	}
	pr_node->count++;
	pr_node->slot[cur_index] = (void*)&data;
	return INSERT_SUCC;
}

//提供給用戶的插入函數：如果當前的樹高不足以存儲index時，需要對樹進行擴展
radix_status radix_tree::radix_tree_insert(u_int index, radix_tree_data& data)
{
	u_int want_level = get_want_level(index);
	cout << index << " want level: " << want_level << endl;

	if(want_level <= r_tree.height) {
		return __radix_tree_insert(index, want_level, data);
	}
	radix_tree_stretch_n(want_level);
	cout << "stretch to " << want_level << " level" << endl;
	return __radix_tree_insert(index, want_level, data);
}

//獲取index所在的內節點
__radix_tree_node* radix_tree::radix_tree_get_node(u_int index)
{
	u_int cur_level = r_tree.height;
	u_int want_level = get_want_level(index);
	u_int cur_index = 0;
	__radix_tree_node *pr_node = r_tree.rnode, *pr_node_tmp = NULL;
	if(want_level <= r_tree.height) {
		while(cur_level > 1) {
			cur_index = get_level_index(index, cur_level);
			pr_node_tmp = (__radix_tree_node*)(pr_node->slot[cur_index]);
			if(pr_node_tmp == NULL) {
				return NULL;
			}
			pr_node = pr_node_tmp;
			--cur_level;
		}
		return pr_node;
	}
	return NULL;
}

//提供給用戶的查詢函數
radix_tree_data* radix_tree::radix_tree_lookup(u_int index)
{
	__radix_tree_node *pr_node = radix_tree_get_node(index);
	if(pr_node == NULL) {
		return NULL;
	}

	u_int cur_index = get_level_index(index, 1);
	radix_tree_data *data = (radix_tree_data*)(pr_node->slot[cur_index]);
	if(data == NULL) {
		return NULL;
	}
	return data;
}

//提供給用戶的刪除函數：只是將index所在的slot置爲NULL，返回存儲的數據，
//當內節點並沒有孩子時，並不刪除內節點
radix_status radix_tree::radix_tree_delete(u_int index, radix_tree_data *& rdata)
{
	__radix_tree_node *pr_node = radix_tree_get_node(index);
	if(pr_node == NULL) {
		return DELETE_FAIL;
	}

	u_int cur_index = get_level_index(index, 1);
	radix_tree_data *data = (radix_tree_data*)(pr_node->slot[cur_index]);
	if(data == NULL) {
		rdata = NULL;
		return DELETE_FAIL;
	}
	rdata = data;
	pr_node->slot[cur_index] = NULL;
	pr_node->count--;
	return DELETE_SUCC;
}

//採用遞歸的方式打印內節點
void radix_tree::radix_tree_print_node(__radix_tree_node *pr_node)
{
	if(pr_node == NULL)
		return;
	if(pr_node->height == 1) {
		for(u_int i = 0; i < RADIX_TREE_SLOTS; ++i) {
			if(pr_node->slot[i]) {
				cout << ((radix_tree_data*)(pr_node->slot[i]))->str << endl;
			}
		}
		return;
	}

	for(u_int i = 0; i < RADIX_TREE_SLOTS; ++i) {
		radix_tree_print_node((__radix_tree_node*)(pr_node->slot[i]));
	}
}

//打印整個樹
void radix_tree::radix_tree_print()
{
	radix_tree_print_node(r_tree.rnode);
}

int main()
{
	radix_tree rdtree;
	u_int index = 0;
	string str;
	ifstream cin("test.txt");
	//test.txt文件內容：
	//5 abc
	//6 def
	//131 mnp

	radix_tree_data data1;
	cin >> index >> str;
	data1.str = str;
	rdtree.radix_tree_insert(index, data1);
	
	radix_tree_data data2;
	cin >> index >> str;
	data2.str = str;
	rdtree.radix_tree_insert(index, data2);
	
	radix_tree_data data3;
	cin >> index >> str;
	data3.str = str;
	rdtree.radix_tree_insert(index, data3);
	
	rdtree.radix_tree_print();
	
	cout << "lookup start..." << endl;
	radix_tree_data *pd = NULL;
	pd = rdtree.radix_tree_lookup(5);
	if(pd)
	    cout << pd->str << endl;
	pd = rdtree.radix_tree_lookup(6);
	if(pd)
	    cout << pd->str << endl;
	pd = rdtree.radix_tree_lookup(131);
	if(pd)
	    cout << pd->str << endl;
	
	return 0;
}