# Augmenting AVL tree to calculate sum of subtree

Suggest a way to augment an AVL tree to support a $$O(\log n)$$ implementation of the function calculateSum(key), which receives a key of a node and returns the sum of its subtree.

I implemented it this way:

sumSubtree(node):
if node != null:
return sumSubtree(node.left) + sumSubtree(node.right) + node.key
return 0

calculateSum(key):
node = Search(key) // assuming I have a search function
return sumSubtree(node)


which solves it in $$O(\log n)$$.

But I read it is possible to maintain the sum during insertion and deletion. And augment an AVL tree this way.

Which solution would be better? Mine, or the other method? Does it matter?

Your solution is traversing the whole subtree on each call to subSubtree so it is actually $$O(|T|) \neq O(\log n)$$ where $$|T|$$ is the length of the subtree. If I query the root of the tree it will go over every node, so complexity is $$O(n)$$.
calculateSum(key):

Search can be implemented in $$O(\log n)$$ and getting the size is $$O(1)$$ since it is stored on the node itself, hence the overall complexity will be $$O(\log n)$$.