public static FlatTree FlattenTree(RandomProjectionTreeNode tree, int leafSize)
{
var nNodes = NumNodes(tree);
var nLeaves = NumLeaves(tree);
// TODO[umap-js]: Verify that sparse code is not relevant...
var hyperplanes = Utils.Range(nNodes).Select(_ => new float[tree.Hyperplane.Length]).ToArray();
var offsets = new float[nNodes];
var children = Utils.Range(nNodes).Select(_ => new[] { -1, -1 }).ToArray();
var indices = Utils.Range(nLeaves).Select(_ => Utils.Range(leafSize).Select(___ => - 1).ToArray()).ToArray();
RecursiveFlatten(tree, hyperplanes, offsets, children, indices, 0, 0);
return(new FlatTree {
Hyperplanes = hyperplanes, Offsets = offsets, Children = children, Indices = indices
});
}
private static (int nodeNum, int leafNum) RecursiveFlatten(RandomProjectionTreeNode tree, float[][] hyperplanes, float[] offsets, int[][] children, int[][] indices, int nodeNum, int leafNum)
{
if (tree.IsLeaf)
{
children[nodeNum][0] = -leafNum;
// TODO[umap-js]: Triple check this operation corresponds to
// indices[leafNum : tree.indices.shape[0]] = tree.indices
tree.Indices.CopyTo(indices[leafNum], 0);
leafNum += 1;
return(nodeNum, leafNum);
}
else
{
hyperplanes[nodeNum] = tree.Hyperplane;
offsets[nodeNum] = tree.Offset;
children[nodeNum][0] = nodeNum + 1;
var oldNodeNum = nodeNum;
var res = RecursiveFlatten(
tree.LeftChild,
hyperplanes,
offsets,
children,
indices,
nodeNum + 1,
leafNum
);
nodeNum = res.nodeNum;
leafNum = res.leafNum;
children[oldNodeNum][1] = nodeNum + 1;
res = RecursiveFlatten(
tree.RightChild,
hyperplanes,
offsets,
children,
indices,
nodeNum + 1,
leafNum
);
return(res.nodeNum, res.leafNum);
}
}
private static int NumLeaves(RandomProjectionTreeNode tree) => tree.IsLeaf ? 1 : (1 + NumLeaves(tree.LeftChild) + NumLeaves(tree.RightChild));
