Advertisement
999ms

Untitled

Jul 23rd, 2020
1,133
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.70 KB | None | 0 0
  1. class BinaryTrie(object):
  2.     def __init__(self, length):
  3.         self.LENGTH = length + 1
  4.         self.MASK = (1 << self.LENGTH) - 1
  5.         self.nodes = [0 for _ in range(64)]
  6.         self.empty_nodes = [63 - i for i in range(63)]
  7.  
  8.     def __sizeof__(self):
  9.         return self.nodes.__sizeof__() + self.empty_nodes.__sizeof__()
  10.  
  11.     def __repr__(self):
  12.         repr = []
  13.         for i in range(len(self.nodes)):
  14.             if i in self.empty_nodes:
  15.                 continue
  16.  
  17.             value = self.nodes[i]
  18.  
  19.             is_leaf = self._get_type(value) == 1
  20.             left = self._get_left(value) or '-'
  21.             right = self._get_right(value) or '-'
  22.  
  23.             if is_leaf:
  24.                 repr.append('%s:leaf' % (i,))
  25.             else:
  26.                 repr.append('%s:%s/%s' % (i, left, right))
  27.  
  28.         return ' '.join(repr)
  29.  
  30.     def _get_left(self, value):
  31.         return (value >> 1) & self.MASK
  32.  
  33.     def _get_right(self, value):
  34.         return value >> (self.LENGTH + 1)
  35.  
  36.     def _get_type(self, value):
  37.         return value & 1
  38.  
  39.     def _create_node(self):
  40.         if len(self.empty_nodes) == 0:
  41.             self.empty_nodes = [2 * len(self.nodes) - 1 - i for i in range(len(self.nodes))]
  42.             self.nodes += [0 for _ in range(len(self.nodes))]
  43.         index = self.empty_nodes.pop()
  44.         return index
  45.  
  46.  
  47.     def _remove_nodes(self, node):
  48.         if node == 0:
  49.             return
  50.        
  51.         self.empty_nodes.append(node)
  52.  
  53.         left_node = self._get_left(self.nodes[node])
  54.         self._remove_nodes(left_node)
  55.  
  56.         right_node = self._get_right(self.nodes[node])
  57.         self._remove_nodes(right_node)
  58.  
  59.         self.nodes[node] = 0
  60.  
  61.     def add(self, value, size=None):
  62.         if size is None:
  63.             size = self.LENGTH
  64.  
  65.         current_node = 0
  66.         for i in range(size):
  67.             bit = size - i - 1
  68.             if self._get_type(self.nodes[current_node]) == 1: # leaf
  69.                 print(current_node, '!')
  70.                 return
  71.             if (value >> bit) & 1 == 0:
  72.                 current_left = self._get_left(self.nodes[current_node])
  73.                 if current_left == 0:
  74.                     current_left = self._create_node()
  75.                     self.nodes[current_node] |= current_left << 1
  76.                 current_node = current_left
  77.             else:
  78.                 current_right = self._get_right(self.nodes[current_node])
  79.                 if current_right == 0:
  80.                     current_right = self._create_node()
  81.                     self.nodes[current_node] |= current_right << (self.LENGTH + 1)
  82.                 current_node = current_right
  83.  
  84.         current_left = self._get_left(self.nodes[current_node]) #delete children if a vertex added
  85.         current_right = self._get_right(self.nodes[current_node])
  86.  
  87.         self._remove_nodes(current_left)
  88.         self._remove_nodes(current_right)
  89.  
  90.         self.nodes[current_node] = 1 #leaf
  91.  
  92.     def __contains__(self, value, size=None):
  93.         if size is None:
  94.             size = self.LENGTH
  95.  
  96.         current_node = 0
  97.         for i in range(size):
  98.             bit = size - i - 1
  99.             if self._get_type(self.nodes[current_node]) == 1:
  100.                 return True
  101.             if (value >> bit) & 1 == 0:
  102.                 current_left = self._get_left(self.nodes[current_node])
  103.                 if current_left == 0:
  104.                     return False
  105.                 current_node = current_left
  106.             else:
  107.                 current_right = self._get_right(self.nodes[current_node])
  108.                 if current_right == 0:
  109.                     return False
  110.                 current_node = current_right
  111.         return self._get_type(self.nodes[current_node]) == 1
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement