You could make your computeWidth()
method completely independent from precomputing the "height" of the tree. (minor nitpick: trees are deep
and not high
).
If performance (and stackframes) is not an issue it's simply a matter of making the function run recursively.
private int computeWidth (TreeNode<?> currentNode, final Map<Integer, Integer> data, int currentLevel) {
if (currentNode == null) {
if (currentLevel == ROOT_LEVEL) {
throw new IllegalArgumentException("Root node mustn't be null");
}
return; //early return, dead end
}
final Integer width = data.get(currentLevel);
data.put(currentLevel, (width == null) ? 1 : width + 1); //increment width for current level
computeWidth(currentNode.left, data, currentLevel + 1);
computeWidth(currentNode.right, data, currentLevel + 1);
return Collections.max(data.values());
}
The idea of this is quite simple. You can store the width of a level in a map, independant of the method you call. You then recursively traverse down the tree, while incrementing a counter you pass with it.
This counter is parallel (or even equivalent) to the depth of the stackframe you are currently in.
The line data.put
is quite hackish, it makes use of boxing and implicit conversions to make a non-existant level in the map not throw NPE's. (null
Integers get converted to 0 when casting to primitive int
)
You'd initially call using following line:
int maxWidth = computeWidth(root, new HashMap<Integer, Integer>(), ROOT_LEVEL);
This way you completely eliminate one n from the big O notation, because you can avoid the precalculation of height, which reduces the timecomplexity from intially \$O(n^3)\$ to just \$O(n^2)\$ in your class because (as mentioned by Pimgdas mentioned by Pimgd) you traverse all input elements to create the tree.
Additionally this approach should also work for any possible Tree
passed into it, given you adjust the internal recursive working of computeWidth
to traverse all branches of a node.