algorithm: suffix_tree

suffix_tree.go

package main

import (
	"fmt"

	"github.com/awalterschulze/gographviz"
)

type Node struct {
	children   map[byte]*Node
	value      string
	compressed bool
}

func build(s string, root *Node) {
	for len(s) != 0 {
		b := s[0]
		if root.children[b] == nil {
			node := NewNode()
			node.value = string(b)
			root.children[b] = node
		}
		root = root.children[b]
		s = s[1:]
	}
}

func buildSuffixTree(s string) *Node {
	root := NewNode()
	root.value = "x"
	for i := 0; i < len(s); i++ {
		build(s[i:], root)
	}
	return root
}

func splitCompressed(root *Node) *Node {
	b := root.value[0]

	node := NewNode()
	node.value = root.value[1:]
	if len(node.value) > 1 {
		node.compressed = true
	}

	root.children[root.value[1]] = node
	root.compressed = false
	root.value = string(b)

	return root
}

func NewNode() *Node {
	node := new(Node)
	node.children = make(map[byte]*Node)
	return node
}

func buildCompressed(s string, root *Node) {
	for len(s) != 0 {
		c := s[0]
		if root.children[c] == nil {
			node := NewNode()
			node.value = s
			node.compressed = true
			root.children[c] = node
			break
		}

		root = root.children[c]
		s = s[1:]
		if root.compressed {
			root = splitCompressed(root)
			continue
		}
	}
}

func buildCompressedSuffixTree(s string) *Node {
	root := new(Node)
	root.children = make(map[byte]*Node)
	root.value = "x"
	for i := 0; i < len(s); i++ {
		buildCompressed(s[i:], root)
	}
	return root
}

func doDump(graph *gographviz.Graph, root *Node, parent string) {
	n += 1
	node := fmt.Sprintf("%s%d", string(root.value), n)
	graph.AddNode("", node, map[string]string{"label": string(root.value)})
	if len(parent) != 0 {
		graph.AddEdge(parent, node, false, nil)
	}
	for _, value := range root.children {
		doDump(graph, value, node)
	}
}

var n int = 0

func dumpSuffixTree(root *Node) {
	graph := gographviz.NewGraph()
	doDump(graph, root, "")
	fmt.Println(graph)
}

func main() {
	// root := buildSuffixTree("ababcadc")
	root := buildCompressedSuffixTree("aabcdabc")
	dumpSuffixTree(root)
}