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.

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

1import java.util.ArrayList;
2import java.util.Arrays;
3import java.util.Iterator;
4import java.util.List;
5import java.util.Map;
6import java.util.Scanner;
7import java.util.function.Function;
8
9class Solution {
10    public static class Node<T> {
11        public T val;
12        public List<Node<T>> children;
13
14        public Node(T val) {
15            this(val, new ArrayList<>());
16        }
17
18        public Node(T val, List<Node<T>> children) {
19            this.val = val;
20            this.children = children;
21        }
22    }
23
24    private float maxAvg = -Float.MAX_VALUE;
25    private Node<Integer> maxRoot;
26
27    // -> (sum, num_nodes)
28    private Map.Entry<Integer, Integer> dfs(Node<Integer> node) {
29        int s = node.val;
30        int n = 1;
31        for (Node<Integer> c : node.children) {
32            Map.Entry<Integer, Integer> e = dfs(c);
33            s += e.getKey();
34            n += e.getValue();
35        }
36        float avg = (float)s / n;
37        if (avg > maxAvg) {
38            maxAvg = avg;
39            maxRoot = node;
40        }
41        return Map.entry(s, n);
42    }
43
44    public static int subtreeMaxAvg(Node<Integer> root) {
45        Solution sol = new Solution();
46        sol.dfs(root);
47        return sol.maxRoot.val;
48    }
49
50    public static <T> Node<T> buildTree(Iterator<String> iter, Function<String, T> f) {
51        String val = iter.next();
52        int num = Integer.parseInt(iter.next());
53        ArrayList<Node<T>> children = new ArrayList<>();
54        for (int i = 0; i < num; i++)
55            children.add(buildTree(iter, f));
56        return new Node<T>(f.apply(val), children);
57    }
58
59    public static List<String> splitWords(String s) {
60        return s.isEmpty() ? List.of() : Arrays.asList(s.split(" "));
61    }
62
63    public static void main(String[] args) {
64        Scanner scanner = new Scanner(System.in);
65        Node<Integer> root = buildTree(splitWords(scanner.nextLine()).iterator(), Integer::parseInt);
66        scanner.close();
67        int res = subtreeMaxAvg(root);
68        System.out.println(res);
69    }
70}
71

Overview

Basic Data Structures

Basic Algorithms

Warm Up

Overview

Sorted Array

Implicitly Sorted Array

Advanced

Introduction

DFS on Tree

Binary Search Tree

Combinatorial Search

Memoization

Pruning

Dedup

Additional Practices

Introduction

BFS on Tree

Introduction

Vanilla BFS

Matrix as Graph

Implicit Graph

Directed Graph / Topological Sort

Weighted Graph

Minimum Spanning Tree

Introduction

Same Direction

Opposite Direction

Sliding Window

Prefix Sum

Cycle Finding

Advanced

Introduction

Top K

Moving Best

Multiple Heaps

Introduction

Sequence

Grid

Dynamic number of subproblems

Partition

Interval

Two Sequences

Game Theory

0-1 Knapsack

Bitmask

Tree

Introduction

Generalized

Binary Search

Disjoint Set Union | Union Find

Trie

Data Structure Design

Segment Tree

Greedy

Monotonic Stack

Line Sweep

Interval

Tree Traversal without Recursion

Composite Patterns

Math

Matrix

Classics

Concepts

Projects

Amazon OA

Microsoft OA

Google OA

Twitter OA

Amazon Online Assessment (OA) - Subtree with Maximum Average

Example 1:

Input:

Output: 12

Explanation

Try it yourself

Explanation

1. Decide on the return value

2. Identify states

Output: `12`