<commit_before>import itertools as it
from collections import deque

from Levenshtein import distance

class Trie:

    _root = None
    _terminals = None
    _collection_count = None

    def __init__(self):
        self._root = Node(element='')
        self._terminals = set()
        self._collection_count = 1

    def add(self, seq, collection_mask):
        node = self._root
        for e in seq:
            node = node.add_or_fetch_child(e)
            node._member_of = node._member_of | collection_mask
        self._terminals.add(node)
          
    def items(self):
        for i in self._terminals:
            yield i.prefix()
     
    def longest_prefix(self, seq):
        node = self._root
        iter = (i for i in seq)

        for i in seq:
            child = node.child_by_element(i)
            if not child:
                break
            node = child
        
        return node

    @staticmethod
    def add_node_active_node_set(node, ed, anset):
        if node not in anset or anset[node] > ed:
            anset[node] = ed 

    @staticmethod
    def calc_active_node_set(node, parent_node_set, sigma):
        active_node_set = {}
        for n, ed in parent_node_set.items():
            if ed + 1 <= sigma:
                Trie.add_node_active_node_set(n, ed + 1, active_node_set)

            for nc in n._child_nodes:
                if not nc._visited:
                    continue
                if nc._element == node._element:
                    Trie.add_node_active_node_set(nc, ed, active_node_set)
                    if ed < sigma and nc != node:
                        for ncc,d in nc.breadth_first(sigma-ed):
                            Trie.add_node_active_node_set(ncc, ed+d, active_node_set)  
                elif ed + 1 <= sigma:
                    Trie.add_node_active_node_set(nc, ed+1, active_node_set)

        return active_node_set 
    
    def trie_search(self, seq, sigma):
        node = self._root
        stack = [None]
        while node:
            if node in self._terminals and node.is_active(seq, sigma):
                yield node
            if not node.can_prune(seq, sigma):
                for child in node:
                    stack.append(child)
            node = stack.pop()

    def _index(self, collection, collection_id):
        for w in collection:
            self.add(w, collection_id)

    def index(self, collection1, collection2=None):
        self._index(collection1, 1)
        if collection2:
            self._collection_count = 2
            self._index(collection2, 2)
   
    def _self_join(self, sigma):
    
        for n in self._root.pre_order():
            n._visited = False
        
        active_node_stack = [{self._root : 0}]
        self._root._visited = True

        traversal_stack = [None] + [c for c in self._root]
        node = traversal_stack.pop()

        while node:
            node._visited = True

            # Calculate this node's active node set
            parent_active_nodes = active_node_stack[-1]
            active_node_set = self.calc_active_node_set(node, parent_active_nodes, sigma)

            # Update active node sets of ancestors
            for ancestor_active_node, cnt in zip(active_node_stack[-1:-sigma-1:-1], it.count(start=1)):
                Trie.add_node_active_node_set(node, cnt, ancestor_active_node)

            # Possibly generate results
            if node in self._terminals:
                for output_candidate in active_node_set:
                    if output_candidate in self._terminals and output_candidate != node:<commit_msg>optimization<commit_after>import itertools as it
from collections import deque

from Levenshtein import distance
from transactionaldict import TransactionalDict as tdict

class Trie:

    _root = None
    _terminals = None
    _collection_count = None

    def __init__(self):
        self._root = Node(element='')
        self._terminals = set()
        self._collection_count = 1

    def add(self, seq, collection_mask):
        node = self._root
        for e in seq:
            node = node.add_or_fetch_child(e)
            node._member_of = node._member_of | collection_mask
        self._terminals.add(node)
          
    def items(self):
        for i in self._terminals:
            yield i.prefix()
     
    def longest_prefix(self, seq):
        node = self._root
        iter = (i for i in seq)

        for i in seq:
            child = node.child_by_element(i)
            if not child:
                break
            node = child
        
        return node

    @staticmethod
    def calc_active_node_set(node, parent_node_set, sigma):
        active_nodes = []
        for n, ed in parent_node_set.items():
            if ed < sigma:
                active_nodes.append((n, ed+1))

            for nc in n._child_nodes:
                if not nc._visited:
                    continue
                if nc._element == node._element:
                    active_nodes.append((nc, ed))
                    if ed < sigma and nc != node:
                        for ncc,d in nc.breadth_first(sigma-ed):
                            active_nodes.append((ncc, ed+d))
                elif ed < sigma:
                    active_nodes.append((nc, ed+1))

        return dict(sorted(active_nodes, key=lambda x: x[1], reverse=True))
    
    def trie_search(self, seq, sigma):
        node = self._root
        stack = [None]
        while node:
            if node in self._terminals and node.is_active(seq, sigma):
                yield node
            if not node.can_prune(seq, sigma):
                for child in node:
                    stack.append(child)
            node = stack.pop()

    def _index(self, collection, collection_id):
        for w in collection:
            self.add(w, collection_id)

    def index(self, collection1, collection2=None):
        self._index(collection1, 1)
        if collection2:
            self._collection_count = 2
            self._index(collection2, 2)
   
    def _self_join(self, sigma):
    
        for n in self._root.pre_order():
            n._visited = False
        
        active_node_stack = [{self._root : 0}]
        self._root._visited = True

        traversal_stack = [None] + [c for c in self._root]
        node = traversal_stack.pop()

        while node:
            node._visited = True

            # Calculate this node's active node set
            parent_active_nodes = active_node_stack[-1]
            active_node_set = self.calc_active_node_set(node, parent_active_nodes, sigma)

            # Update active node sets of ancestors
            for ancestor_active_node, cnt in zip(active_node_stack[-1:-sigma-1:-1], it.count(start=1)):
                if not node in ancestor_active_node or cnt < ancestor_active_node[node]:
                    ancestor_active_node[node] = cnt

            # Possibly generate results
            if node in self._terminals:
                for output_candidate in active_node_set:
                    if output_candidate in self._terminals and output_candidate != node:
{'fraction_non_alphanumeric': 0.07856941706561532, 'fraction_numerical': 0.005209800056322163, 'mean_word_length': 1.7183314198239572, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '868674', 'n_tokens_mistral': 2097, 'n_tokens_neox': 1998, 'n_words': 536}