Amazon Online Assessment (OA) - Subtree with Maximum Average

Given a tree, find the subtree with the maximum average. Return the subtree's root's value. Note that the tree could have arbitrarily many children.

Example 1:

Input:

Output: 12

Explanation

The sum of each subtree:

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

Try it yourself

Explanation

Prereq: DFS on Tree, DFS on ternary tree

This is a pretty straightforward question - we traverse the tree and the sum and number of nodes in the subtree to the parent.

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 ...