代码

#include "router.h"
#include <stdint.h>
#include <stdlib.h>
#include <arpa/inet.h>
#include <vector>
#include <iostream>
using namespace std;

/*
  RoutingTable Entry 的定义如下：
  typedef struct {
    uint32_t addr; // 大端序，IPv4 地址
    uint32_t len; // 小端序，前缀长度
    uint32_t if_index; // 小端序，出端口编号
    uint32_t nexthop; // 大端序，下一跳的 IPv4 地址
  } RoutingTableEntry;

  约定 addr 和 nexthop 以 **大端序** 存储。
  这意味着 1.2.3.4 对应 0x04030201 而不是 0x01020304。
  保证 addr 仅最低 len 位可能出现非零。
  当 nexthop 为零时这是一条直连路由。
  你可以在全局变量中把路由表以一定的数据结构格式保存下来。
*/

struct Node {
  bool left = true;
  // int entryNum = 0;
  // RoutingTableEntry * entry = nullptr;
  std::vector<RoutingTableEntry> entry;
  Node * parent = nullptr;
  Node * lNode = nullptr, * rNode = nullptr;
};

void copyEntry(RoutingTableEntry * entry, RoutingTableEntry oldEntry) {
  entry -> addr = oldEntry.addr;
  entry -> len = oldEntry.len;
  entry -> nexthop = oldEntry.nexthop;
}

void copyNodeEntry(RoutingTableEntry * entry, Node * node, uint32_t addr) {
  copyEntry(entry, node -> entry[0]);
  for (int i = 0; i < node -> entry.size(); i ++) {
    if (node -> entry[i].addr == addr) {
      copyEntry(entry, node -> entry[i]);
      break;
    }
  }
}

class RoutingTrie {
  static Node * root;
public:
  static RoutingTableEntry * search(uint32_t addr) {
    uint32_t newAddr = ntohl(addr);
    Node * node = root;
    // RoutingTableEntry * entry = root -> entry;
    RoutingTableEntry * entry = nullptr;
    if (!node -> entry.empty()) {
      entry = new RoutingTableEntry;
      copyNodeEntry(entry, node, addr);
    }
    for (int i = 0; i < 32; i ++) {
      int bit = newAddr >> 31;
      node = (bit == 0) ? node -> lNode : node -> rNode;
      if (node != nullptr) {
        // entry = (node -> entry == nullptr) ? entry : node -> entry; 
        if (!node -> entry.empty()) {
          entry = new RoutingTableEntry;
          copyNodeEntry(entry, node, addr);
        }
      } else {
        break;
      }
      newAddr = newAddr << 1;
    }
    return entry;
  }

  static void insertNode(RoutingTableEntry entry) {
    uint32_t addr = ntohl(entry.addr);
    int len = entry.len;
    Node * node = root;

    for (int i = 0; i < len; i ++) {
      int bit = addr >> 31;
      if (bit == 0) {
        if (node -> lNode == nullptr) {
          node -> lNode = new Node;
          node -> lNode -> parent = node;
        }
        node = node -> lNode;
      } else {
        if (node -> rNode == nullptr) {
          node -> rNode = new Node;
          node -> rNode -> parent = node;
          node -> rNode -> left = false;
        }
        node = node -> rNode;
      }
      addr = addr << 1;
    }

    // if (node -> entry != nullptr) {
    //   if (node -> entry -> addr == entry.addr) {
    //     copyEntry(node -> entry, entry);
    //   } else {
    //     node -> entryNum ++;
    //   }
    // } else {
    //   node -> entry = new RoutingTableEntry;
    //   copyEntry(node -> entry, entry);
    //   node -> entryNum ++;
    // }
    bool found = false;
    for (int i = 0; i < node -> entry.size(); i ++) {
      if (node -> entry[i].addr == entry.addr) {
        node -> entry[i].nexthop = entry.nexthop;
        found = true;
        break;
      }
    }
    if (!found) {
      node -> entry.push_back(entry);
    }
  }
  
  static void deleteNode(RoutingTableEntry entry) {
    uint32_t addr = ntohl(entry.addr), len = entry.len;
    Node * node = root;
    while (len != 0) {
      int bit = addr >> 31;
      node = (bit == 0) ? node -> lNode : node -> rNode;
      if (node != nullptr) {
        addr = addr << 1;
        len --;
      } else {
        break;
      }
    }
    if (len == 0) { //find node
      // if (node -> entry != nullptr)
      //   node -> entryNum --;
      // if (node -> entryNum == 0) {
      //   while (node != root && node -> lNode == nullptr && node -> rNode == nullptr && node -> entryNum == 0) {
      //     bool leftNode = node -> left;
      //     node = node -> parent;
      //     if (leftNode) {
      //       delete node -> lNode;
      //       node -> lNode = nullptr;
      //     } else {
      //       delete node -> rNode;
      //       node -> rNode = nullptr;
      //     }
      //   }
      // }
      for (int i = 0; i < node -> entry.size(); i ++) {
        if (node -> entry[i].addr == entry.addr) {
          node -> entry.erase(node -> entry.begin() + i);
          break;
        }
      }
      if (node -> entry.empty()) {
        while (node != root && node -> lNode == nullptr && node -> rNode == nullptr && node -> entry.empty()) {
          bool leftNode = node -> left;
          node = node -> parent;
          if (leftNode) {
            delete node -> lNode;
            node -> lNode = nullptr;
          } else {
            delete node -> rNode;
            node -> rNode = nullptr;
          }
        }
      }
    }
  }

};

Node * RoutingTrie::root = new Node;

void init(int n, int q, const RoutingTableEntry *a) {
  for (int i = 0; i < n; i ++) {
    RoutingTrie::insertNode(a[i]);
  }
}

unsigned query(unsigned addr) {
  RoutingTableEntry * entry = RoutingTrie::search(addr);
  if (entry != nullptr) {
    return entry -> nexthop;
  }
  return 0;
}

评测结果