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();
}



Go: package main

import "fmt"

func main() {
    vis(buildTree("banana$"))
}

type tree []node

type node struct {
    sub string 
    ch  []int  
}

func buildTree(s string) tree {
    t := tree{node{}} 
    for i := range s {
        t = t.addSuffix(s[i:])
    }
    return t
}

func (t tree) addSuffix(suf string) tree {
    n := 0
    for i := 0; i < len(suf); {
        b := suf[i]
        ch := t[n].ch
        var x2, n2 int
        for ; ; x2++ {
            if x2 == len(ch) {
                
                n2 = len(t)
                t = append(t, node{sub: suf[i:]})
                t[n].ch = append(t[n].ch, n2)
                return t
            }
            n2 = ch[x2]
            if t[n2].sub[0] == b {
                break
            }
        }
        
        sub2 := t[n2].sub
        j := 0
        for ; j < len(sub2); j++ {
            if suf[i+j] != sub2[j] {
                
                n3 := n2
                
                n2 = len(t)
                t = append(t, node{sub2[:j], []int{n3}})
                t[n3].sub = sub2[j:] 
                t[n].ch[x2] = n2
                break 
            }
        }
        i += j 
        n = n2 
    }
    return t
}

func vis(t tree) {
    if len(t) == 0 {
        fmt.Println("<empty>")
        return
    }
    var f func(int, string)
    f = func(n int, pre string) {
        children := t[n].ch
        if len(children) == 0 {
            fmt.Println("╴", t[n].sub)
            return
        }
        fmt.Println("┐", t[n].sub)
        last := len(children) - 1
        for _, ch := range children[:last] {
            fmt.Print(pre, "├─")
            f(ch, pre+"│ ")
        }
        fmt.Print(pre, "└─")
        f(children[last], pre+"  ")
    }
    f(0, "")
}