~/problems / Trees / Tree DP

Pruning a diseased fruit tree

easy ~15 min

A fruit tree has n junctions numbered 0 .. n-1; junction 0 is the trunk and parent[i] is the junction directly below i (with parent[0] == -1). Each junction carries fruit worth value[i], which is negative where the wood is diseased (keeping it costs you more in treatment than its fruit brings).

A gardener can saw through the branch between any junction i != 0 and its parent; that removes i and everything above it, and each cut costs cut_cost. The trunk always stays.

Implement best_harvest(parent: list[int], value: list[int], cut_cost: int) -> int: the largest possible value of the junctions still attached to the trunk, minus the cost of all the cuts made.

#        0 (3)
#       /     \
#   1 (-2)    2 (4)
#   /    \
# 3 (5)  4 (-6)
best_harvest([-1, 0, 0, 1, 1], [3, -2, 4, 5, -6], 1)  # 9   cut only junction 4: 3 - 2 + 4 + 5 - 1
best_harvest([-1, 0], [1, -1], 5)                     # 0   a cut costs more than it saves

Constraints: 1 <= n <= 100_000, -1000 <= value[i] <= 1000, 0 <= cut_cost <= 1000. The junctions are not numbered in any particular order, and the tree may be one long chain, so a plain recursive DFS will hit Python's recursion limit: traverse iteratively.

Show hint

let keep[v] be the best result for the part of the tree above v, given that v itself stays. Each child c either stays (worth keep[c]) or is cut (worth -cut_cost), so keep[v] = value[v] + sum of max(keep[c], -cut_cost) over its children; the answer is keep[0]. Process children before parents, e.g. in reverse BFS order.

Topic: Tree DP. Post-order: each node returns a small tuple of states to its parent.

Read the visual guide
0:00
Ctrl ' run · Ctrl ↵ submit
esc