算法题 hihoCoder 1289

2016-04-07 作者: Zhenrui Chen 标签: 编程
原文链接: http://zhenruichen.com/2016/04/07/hihoCode-1289.html

最开始实现了一个暴力遍历的算法，需要查询比对 M*N 次。算法逻辑上没有问题，但是提交之后显示 Time Limit Exceeded. 之后实现了基于前缀树 trie 的搜索算法，提交后结果为 Accepted.

算法思路

把 ipv4 地址转换成32个比特，建立一个32层的二叉树。根据过滤规则，在每个结点标注状态，0表示不确定，1表示接受，2表示拒绝。左子树表示下一比特为0，右子树表示下一比特为1.
在读入每一条规则时，把 trie 中对应的结点标记为规则中描述的状态。如果 trie 中没有对应的结点，先创建新结点再标注状态。同时，添加当前结点的优先级，越早添加的结点优先级越高。
读入一条需要判断的 ipv4 地址时，在 trie 中查找目标结点的状态，相应地返回 true 或者 false. 如果没有对应的结点，返回 true.

代码实现

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
using namespace std;

vector<bool> getBinary(const vector<int> &reqIP) {
    vector<bool> bits(32, false);
    for (int i = 0; i < 4; i++) {
        int num = reqIP[i];
        for (int j = 0; j < 8; j++) {
            bits[8 * i + 7 - j] = (num % 2 == 1);
            num = num >> 1;
        }
    }
    return bits;
}

class TreeNode {
public:
    // 0: not sure, 1:allow, 2:deny
    int type;
    int priority;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int type, int priority) {
        this->type = type;
        this->priority = priority;
        this->left = NULL;
        this->right = NULL;
    }
};

bool isAllowed(TreeNode *root, const vector<bool> &bitsReqIP) {
    TreeNode *head = root;
    bool state = true;
    int maxPriority = -1;
    for (int i = 0; i < 32; i++) {
        if (head->type > 0 && head->priority > maxPriority) {
            maxPriority = head->priority;
            state = (head->type == 1);
        }
        if (!bitsReqIP[i]) {
            head = head->left;
        } else {
            head = head->right;
        }
        if (head == NULL) {
            break;
        }
    }
    if (head != NULL && head->priority > maxPriority) {
        state = (head->type != 2);
    }
    return state;
}

int main() {

    int N, M;
    cin >> N >> M;

    TreeNode * root = new TreeNode(0, 0);
    for (int i = 0; i < N; i++) {
        string allow, slash;
        int mask = -1;
        char c;
        vector<int> ip(4, 0);

        cin >> allow >> ip[0] >> c >> ip[1] >> c >> ip[2] >> c >> slash;
        auto loc = slash.find('/', 0);
        if ( loc == string::npos) {
            ip[3] = stoi(slash, 0, 10);
        } else {
            string strIP(slash, 0, loc + 1);
            string strMask(slash, loc + 1);
            ip[3] = stoi(strIP, 0, 10);
            mask = stoi(strMask, 0, 10);
        }

        vector<bool> bits = getBinary(ip);
        TreeNode * head = root;
        int range = (mask < 0) ? 32 : mask;
        for (int m = 0; m < range; m++) {
            if (!bits[m]) {
                if (head->left == NULL) {
                    head->left = new TreeNode(0, 0);
                }
                head = head->left;
            } else {
                if (head->right == NULL) {
                    head->right = new TreeNode(0, 0);
                }
                head = head->right;
            }
        }
        if (head->priority < N - i) {
            head->priority = N - i;
            if (allow == "deny") {
                head->type = 2;
            } else {
                head->type = 1;
            }
        }

    }

    for (int i = 0; i < M; i++) {
        char c;
        vector<int> ip(4, -1);
        cin >> ip[0] >> c >> ip[1] >> c >> ip[2] >> c >> ip[3];

        vector<bool> bitsReqIP = getBinary(ip);
        if (isAllowed(root, bitsReqIP)) {
            cout << "YES" << endl;
        } else {
            cout << "NO" << endl;
        }
    }

    return 0;
}