func splitExamples(t *pb.TreeNode, e Examples) (left Examples, right Examples) {
by(func(e1, e2 *pb.Example) bool {
return e1.Features[t.GetFeature()] < e2.Features[t.GetFeature()]
}).Sort(e)
splitIndex := 0
for i, ex := range e {
splitIndex = i
if ex.Features[t.GetFeature()] > t.GetSplitValue() {
break
}
}
left, right = e[:splitIndex], e[splitIndex:]
return
}
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)
}
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")
}