61 lines
1.9 KiB
Go
61 lines
1.9 KiB
Go
package decisionTree
|
|
|
|
type GraphPart struct {
|
|
Label string
|
|
Children map[string]GraphPart `json:"Children,omitempty"`
|
|
}
|
|
|
|
// Every question is a non-leaf vertex. Every answer is an edge. The root question is the one which is never a target
|
|
// of an answer. We'll start by finding it by using that definition.
|
|
var vertices = make(map[string]Vertex)
|
|
var edges []Edge
|
|
|
|
func GraphFromDBRows(vs []Vertex, es []Edge) GraphPart {
|
|
// As our graph is a simple tree (as opposed to polytree etc), there is always a single root node. It's not strictly
|
|
// necessary to find that node up front, but it does help make this example simpler.
|
|
rootQuestion := getRootVertex(vs, es)
|
|
|
|
// Since building the graph is a recursive operation, we can keep the stack a bit lighter (and our function signatures
|
|
// simpler) by storing the questions and edges in package variables.
|
|
edges = es
|
|
|
|
// Storing questions as a map for easy lookup by ID later
|
|
for _, v := range vs {
|
|
vertices[v.ID] = v
|
|
}
|
|
|
|
// Now that we have the root question, we can recursively build the rest of the tree
|
|
graph := GraphPart{Label: rootQuestion.Value, Children: make(map[string]GraphPart)}
|
|
buildChildren(&graph, rootQuestion)
|
|
|
|
// And we can return the graph, which provides references to the full graph
|
|
return graph
|
|
}
|
|
|
|
func buildChildren(n *GraphPart, v Vertex) {
|
|
for _, e := range edges {
|
|
if e.SourceID == v.ID {
|
|
child := GraphPart{Label: vertices[e.TargetID].Value, Children: make(map[string]GraphPart)}
|
|
buildChildren(&child, vertices[e.TargetID])
|
|
n.Children[e.Value] = child
|
|
}
|
|
}
|
|
}
|
|
|
|
func getRootVertex(vs []Vertex, es []Edge) Vertex {
|
|
targetIDs := map[string]struct{}{}
|
|
|
|
for _, a := range es {
|
|
targetIDs[a.TargetID] = struct{}{}
|
|
}
|
|
|
|
for _, v := range vs {
|
|
if _, present := targetIDs[v.ID]; !present {
|
|
return v
|
|
}
|
|
}
|
|
|
|
// A lazy panic. In the real world we should do nice error handling.
|
|
panic("couldn't find the root question")
|
|
}
|