initial commit

decisionTree/graph.go

@@ -0,0 +1,62 @@
+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")
+}