Amazon Online Assessment (OA) - Subtree with Maximum Average

Given a tree, find the subtree with the maximum average value. Return the subtree's root's value. Note that the tree could have an arbitrary number of children.

Example 1:

Input:

Output: 12

Explanation

The sum of each subtree:

The subtree max average is 9.3, and the root of the subtree with max average is the node with value 12. So we return 12.

Try it yourself

Explanation

Prereq: DFS on Tree, DFS on ternary tree

We traverse the tree with DFS and keep track of the sum of the subtree and the number of its children nodes.

1. Decide on the return value

The problem asks for the root node with the max average, so we update the max average value and its corresponding root node value as we go.

2. Identify states

To calculate the average of the subtree at each node, we need the subtree's sum and the number of children nodes. We can carry this as a state as we traverse the tree.

Having decided on the state and return value we can now write the DFS.

1
+
from math import inf
2
+
from typing import Tuple
3
+
1
4
class Node:
2
5
    def __init__(self, val, children=None):
3
6
        if children is None:
Expand 2 lines ...
6
9
        self.children = children
7
10
8
11
def subtree_max_avg(root: Node) -> int:
9
-
    # WRITE YOUR BRILLIANT CODE HERE
12
+
    # (max_avg, root)
10
-
    return 0
13
+
    res = (-inf, None)
11
14
12
15
16
+
    # -> (sum, num_nodes)
17
+
    def dfs(node: Node) -> Tuple[int, int]:
18
+
        nonlocal res
19
+
        rec = [dfs(c) for c in node.children if c]
20
+
        s = node.val + sum(t[0] for t in rec)
21
+
        n = 1 + sum(t[1] for t in rec)
22
+
        res = max(res, (s / n, node))
23
+
        return s, n
24
+
25
+
    dfs(root)
26
+
    return res[1].val
27
+
13
28
def build_tree(nodes, f):
14
29
    val = next(nodes)
15
30
    num = int(next(nodes))
16
31
    children = [build_tree(nodes, f) for _ in range(num)]
17
32
    return Node(f(val), children)
Expand 5 lines ...