func printNode(node *pb.TreeNode, c *codeWriter) { if node.GetLeft() == nil && node.GetRight() == nil { c.WriteString(fmt.Sprintf("return %v;\n", node.GetLeafValue())) return } c.WriteString(fmt.Sprintf("if (%v(f[%v] < %v)) {\n", getAnnotation(node), node.GetFeature(), node.GetSplitValue())) { c.indentLevel++ printNode(node.GetLeft(), c) c.indentLevel-- } c.WriteString("} else {\n") { c.indentLevel++ printNode(node.GetRight(), c) c.indentLevel-- } c.WriteString("}\n") }
func flattenTree(f *fastTreeEvaluator, current *pb.TreeNode, currentIndex int) { glog.Infof("Flattening tree at index %v", currentIndex) if isLeaf(current) { f.nodes[currentIndex] = flatNode{ value: current.GetLeafValue(), feature: leafFeatureID, } return } // append child nodes // since we push on N + 2 elements, we want index N + 1, hence len(f.nodes) leftChild := len(f.nodes) f.nodes = append(f.nodes, flatNode{}, flatNode{}) f.nodes[currentIndex] = flatNode{ value: current.GetSplitValue(), feature: current.GetFeature(), leftChild: leftChild, } flattenTree(f, current.GetLeft(), leftChild) flattenTree(f, current.GetRight(), leftChild+1) }