Task: Suffix tree


C++: #include <functional>
#include <iostream>
#include <vector>

struct Node {
    std::string sub = "";   
    std::vector<int> ch;    

    Node() {
        
    }

    Node(const std::string& sub, std::initializer_list<int> children) : sub(sub) {
        ch.insert(ch.end(), children);
    }
};

struct SuffixTree {
    std::vector<Node> nodes;

    SuffixTree(const std::string& str) {
        nodes.push_back(Node{});
        for (size_t i = 0; i < str.length(); i++) {
            addSuffix(str.substr(i));
        }
    }

    void visualize() {
        if (nodes.size() == 0) {
            std::cout << "<empty>\n";
            return;
        }

        std::function<void(int, const std::string&)> f;
        f = [&](int n, const std::string & pre) {
            auto children = nodes[n].ch;
            if (children.size() == 0) {
                std::cout << "- " << nodes[n].sub << '\n';
                return;
            }
            std::cout << "+ " << nodes[n].sub << '\n';

            auto it = std::begin(children);
            if (it != std::end(children)) do {
                if (std::next(it) == std::end(children)) break;
                std::cout << pre << "+-";
                f(*it, pre + "| ");
                it = std::next(it);
            } while (true);

            std::cout << pre << "+-";
            f(children[children.size() - 1], pre + "  ");
        };

        f(0, "");
    }

private:
    void addSuffix(const std::string & suf) {
        int n = 0;
        size_t i = 0;
        while (i < suf.length()) {
            char b = suf[i];
            int x2 = 0;
            int n2;
            while (true) {
                auto children = nodes[n].ch;
                if (x2 == children.size()) {
                    
                    n2 = nodes.size();
                    nodes.push_back(Node(suf.substr(i), {}));
                    nodes[n].ch.push_back(n2);
                    return;
                }
                n2 = children[x2];
                if (nodes[n2].sub[0] == b) {
                    break;
                }
                x2++;
            }
            
            auto sub2 = nodes[n2].sub;
            size_t j = 0;
            while (j < sub2.size()) {
                if (suf[i + j] != sub2[j]) {
                    
                    auto n3 = n2;
                    
                    n2 = nodes.size();
                    nodes.push_back(Node(sub2.substr(0, j), { n3 }));
                    nodes[n3].sub = sub2.substr(j); 
                    nodes[n].ch[x2] = n2;
                    break; 
                }
                j++;
            }
            i += j; 
            n = n2; 
        }
    }
};

int main() {
    SuffixTree("banana$").visualize();
}






Python: class Node:
    def __init__(self, sub="", children=None):
        self.sub = sub
        self.ch = children or []

class SuffixTree:
    def __init__(self, str):
        self.nodes = [Node()]
        for i in range(len(str)):
            self.addSuffix(str[i:])

    def addSuffix(self, suf):
        n = 0
        i = 0
        while i < len(suf):
            b = suf[i]
            x2 = 0
            while True:
                children = self.nodes[n].ch
                if x2 == len(children):
                    
                    n2 = len(self.nodes)
                    self.nodes.append(Node(suf[i:], []))
                    self.nodes[n].ch.append(n2)
                    return
                n2 = children[x2]
                if self.nodes[n2].sub[0] == b:
                    break
                x2 = x2 + 1

            
            sub2 = self.nodes[n2].sub
            j = 0
            while j < len(sub2):
                if suf[i + j] != sub2[j]:
                    
                    n3 = n2
                    
                    n2 = len(self.nodes)
                    self.nodes.append(Node(sub2[:j], [n3]))
                    self.nodes[n3].sub = sub2[j:] 
                    self.nodes[n].ch[x2] = n2
                    break 
                j = j + 1
            i = i + j   
            n = n2      

    def visualize(self):
        if len(self.nodes) == 0:
            print "<empty>"
            return

        def f(n, pre):
            children = self.nodes[n].ch
            if len(children) == 0:
                print "--", self.nodes[n].sub
                return
            print "+-", self.nodes[n].sub
            for c in children[:-1]:
                print pre, "+-",
                f(c, pre + " | ")
            print pre, "+-",
            f(children[-1], pre + "  ")

        f(0, "")

SuffixTree("banana$").visualize()