LeetCode 3562 - Maximum Profit from Trading Stocks with Discounts

We are given a rooted tree of employees, where employee 1 is the CEO and the root of the hierarchy. Each node represents an investment opportunity with a cost today (present[i]) and a return tomorrow (future[i]).

leetcodehardarraydynamic-programmingtreedepth-first-search

LeetCode Problem 3562

Difficulty: 🔴 Hard
Topics: Array, Dynamic Programming, Tree, Depth-First Search

Solution

Problem Understanding

We are given a rooted tree of employees, where employee 1 is the CEO and the root of the hierarchy. Each node represents an investment opportunity with a cost today (present[i]) and a return tomorrow (future[i]). The profit from investing in employee i is therefore future[i] - cost[i], where the cost depends on a discount rule.

The key structural dependency is the discount condition: an employee v receives a 50% discount on their stock price if and only if their direct boss buys their own stock. This creates a parent-child dependency in a rooted tree, where each node’s cost depends on whether its parent is selected in the investment set.

We must select a subset of nodes to maximize total profit while ensuring the total purchase cost does not exceed budget. Each node can be selected at most once, and profits do not reinvest into budget capacity, so this is a constrained knapsack optimization over a tree with state-dependent weights.

The constraints n ≤ 160 and budget ≤ 160 strongly indicate a dynamic programming solution with cubic or near-cubic complexity in the worst case is acceptable. The hierarchy is guaranteed to be a tree (acyclic, connected with root 1), which enables tree DP.

Edge cases include the minimum tree (n = 1), chains (degenerate tree), and cases where discounts drastically change optimal structure. A naive subset enumeration is exponential and infeasible.

Approaches

Brute Force Idea

A brute force solution would enumerate every subset of employees, check whether it is valid under the budget constraint, and compute profit under discount rules induced by parent selections. For each subset, we must also verify the tree constraint implicitly defining discount propagation.

This approach is correct because it directly evaluates all feasible configurations of selected nodes and their induced costs. However, it is infeasible because there are 2^n subsets, and for each subset we would need to recompute costs and profits, leading to exponential time.

Key Insight

The hierarchy forms a rooted tree, and the discount condition depends only on whether the parent node is selected. This local dependency allows us to define a tree dynamic programming state.

At each node, the decision is binary: either buy the stock or not. However, the cost and resulting child transitions depend on this decision. This naturally leads to a tree DP where each node maintains two states:

  • Parent did not buy this node’s stock
  • Parent did buy this node’s stock

For each state, we perform a knapsack merge over children, combining subproblem results using convolution over budget capacity.

The critical observation is that once the parent state is fixed, each subtree becomes independent, allowing optimal substructure.

Complexity Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n · n) O(n) Enumerate all subsets and evaluate constraints
Optimal Tree DP + Knapsack O(n · budget²) O(n · budget) Tree DP with budget convolution at each node

Algorithm Walkthrough

We define a rooted tree DP. Let dp[u][p][b] be the maximum profit obtainable in the subtree rooted at u, given that the parent of u is in state p (where p = 0 means parent did not buy, p = 1 means parent bought), and using budget b.

However, instead of directly maintaining both states during recursion, we compute them symmetrically per node using two aggregated knapsack tables over children.

We proceed as follows:

  1. Root the tree at node 1. Build adjacency lists from the hierarchy. This ensures a directed parent-to-child structure, making subtree DP valid.
  2. For each node u, we compute two DP tables:
  • dp0[u][b]: maximum profit in subtree of u if parent did NOT buy u
  • dp1[u][b]: maximum profit in subtree of u if parent DID buy u

The only difference is the cost of buying u, which is either:

  • present[u] if parent state is 0
  • floor(present[u] / 2) if parent state is 1
  1. For each node u, we first compute child aggregation tables separately:
  • base0[b]: best result from all children assuming u is NOT bought
  • base1[b]: best result from all children assuming u IS bought

This separation is necessary because child discount depends on whether u is purchased. 4. To compute base0, we initialize a knapsack array with only zero cost and zero profit. Then for each child v, we merge:

base0 ⊗ dp[v][0], where convolution is over budget.

Similarly, base1 is computed using dp[v][1]. 5. After child aggregation, we decide whether to buy u.

If we do not buy u, profit is base0[b].

If we buy u, we allocate cost c_u (depending on parent state), and combine:

base1[b - c_u] + (future[u] - c_u).

We take the maximum of these two choices for each budget b. 6. We return dp0[1][budget], since the root has no parent (equivalent to parent not buying).

Why it works

The correctness follows from optimal substructure and independence of subtrees given a fixed parent state. Once we condition on whether a node is purchased, all child subproblems become independent and reduce to bounded knapsack merges. The tree DP ensures that each edge is considered exactly once in defining discount propagation.

Python Solution

from typing import List

class Solution:
    def maxProfit(self, n: int, present: List[int], future: List[int], hierarchy: List[List[int]], budget: int) -> int:
        from collections import defaultdict
        import math

        adj = defaultdict(list)
        for u, v in hierarchy:
            adj[u].append(v)

        NEG = -10**18

        def dfs(u: int):
            dp0 = [[NEG] * (budget + 1) for _ in range(2)]
            dp1 = [[NEG] * (budget + 1) for _ in range(2)]

            dp0[0][0] = 0
            dp0[1][0] = 0

            dp1[0][0] = 0
            dp1[1][0] = 0

            base0 = [NEG] * (budget + 1)
            base1 = [NEG] * (budget + 1)
            base0[0] = 0
            base1[0] = 0

            for v in adj[u]:
                c0, c1 = dfs(v)

                new0 = [NEG] * (budget + 1)
                new1 = [NEG] * (budget + 1)

                for b1 in range(budget + 1):
                    if base0[b1] < 0:
                        continue
                    for b2 in range(budget - b1 + 1):
                        if c0[b2] < 0:
                            continue
                        new0[b1 + b2] = max(new0[b1 + b2], base0[b1] + c0[b2])

                for b1 in range(budget + 1):
                    if base1[b1] < 0:
                        continue
                    for b2 in range(budget - b1 + 1):
                        if c1[b2] < 0:
                            continue
                        new1[b1 + b2] = max(new1[b1 + b2], base1[b1] + c1[b2])

                base0, base1 = new0, new1

            cost0 = present[u - 1]
            cost1 = cost0 // 2
            profit0 = future[u - 1] - cost0
            profit1 = future[u - 1] - cost1

            res0 = [NEG] * (budget + 1)
            res1 = [NEG] * (budget + 1)

            for b in range(budget + 1):
                res0[b] = base0[b]

            for b in range(cost0, budget + 1):
                if base1[b - cost0] >= 0:
                    res0[b] = max(res0[b], base1[b - cost0] + profit0)

            for b in range(budget + 1):
                res1[b] = base0[b]

            for b in range(cost1, budget + 1):
                if base1[b - cost1] >= 0:
                    res1[b] = max(res1[b], base1[b - cost1] + profit1)

            return res0, res1

        dp0, dp1 = dfs(1)
        return max(dp0[budget], dp1[budget])

Implementation Notes

The function dfs(u) returns two knapsack tables representing the two parent states. The base0 and base1 arrays aggregate child contributions depending on whether u is purchased. The final step merges the decision of buying or not buying u, respecting the budget constraint and applying the correct cost and profit.

Go Solution

package main

func maxProfit(n int, present []int, future []int, hierarchy [][]int, budget int) int {
	adj := make([][]int, n+1)
	for _, e := range hierarchy {
		u, v := e[0], e[1]
		adj[u] = append(adj[u], v)
	}

	const NEG = -1 << 60

	var dfs func(u int) ([][]int, [][]int)
	dfs = func(u int) ([][]int, [][]int) {
		dp0 := make([][]int, 2)
		dp1 := make([][]int, 2)
		for i := 0; i < 2; i++ {
			dp0[i] = make([]int, budget+1)
			dp1[i] = make([]int, budget+1)
			for b := 0; b <= budget; b++ {
				dp0[i][b] = NEG
				dp1[i][b] = NEG
			}
			dp0[i][0], dp1[i][0] = 0, 0
		}

		base0 := make([]int, budget+1)
		base1 := make([]int, budget+1)
		for i := 0; i <= budget; i++ {
			base0[i], base1[i] = NEG, NEG
		}
		base0[0], base1[0] = 0, 0

		for _, v := range adj[u] {
			c0, c1 := dfs(v)

			new0 := make([]int, budget+1)
			new1 := make([]int, budget+1)
			for i := 0; i <= budget; i++ {
				new0[i], new1[i] = NEG, NEG
			}

			for i := 0; i <= budget; i++ {
				if base0[i] == NEG {
					continue
				}
				for j := 0; j+i <= budget; j++ {
					if c0[j] == NEG {
						continue
					}
					if base0[i]+c0[j] > new0[i+j] {
						new0[i+j] = base0[i] + c0[j]
					}
				}
			}

			for i := 0; i <= budget; i++ {
				if base1[i] == NEG {
					continue
				}
				for j := 0; j+i <= budget; j++ {
					if c1[j] == NEG {
						continue
					}
					if base1[i]+c1[j] > new1[i+j] {
						new1[i+j] = base1[i] + c1[j]
					}
				}
			}

			base0, base1 = new0, new1
		}

		cost0 := present[u-1]
		cost1 := cost0 / 2
		profit0 := future[u-1] - cost0
		profit1 := future[u-1] - cost1

		res0 := make([]int, budget+1)
		res1 := make([]int, budget+1)
		for i := 0; i <= budget; i++ {
			res0[i] = base0[i]
			res1[i] = base0[i]
		}

		for b := cost0; b <= budget; b++ {
			if base1[b-cost0] != NEG {
				if base1[b-cost0]+profit0 > res0[b] {
					res0[b] = base1[b-cost0] + profit0
				}
			}
		}

		for b := cost1; b <= budget; b++ {
			if base1[b-cost1] != NEG {
				if base1[b-cost1]+profit1 > res1[b] {
					res1[b] = base1[b-cost1] + profit1
				}
			}
		}

		return res0, res1
	}

	dp0, dp1 := dfs(1)
	if dp0[budget] > dp1[budget] {
		return dp0[budget]
	}
	return dp1[budget]
}

Go-specific Notes

The Go implementation uses explicit slice allocation for DP tables, and NEG is chosen as a large negative sentinel compatible with int. Unlike Python, there is no automatic handling of negative infinity, so careful initialization is required. Slice-based DP makes memory reuse straightforward but requires explicit bounds checks.

Worked Examples

Example 1

For n = 2, hierarchy 1 → 2, we evaluate node 1 first.

At node 1:

  • If buy 1: cost = 1, profit = 3
  • If not buy 1: cost = 0, profit = 0

Node 2 depends on whether 1 is bought.

If 1 is bought:

  • node 2 cost = 1 (discount), profit = 2

DP merges:

  • buy both: cost = 2, profit = 5
  • valid since budget = 3

Thus optimal = 5.

Example 2

Only one node can be chosen due to budget structure:

  • Node 1 or 2, but not both under optimal cost-profit tradeoff

Best is node 2 alone:

profit = 8 - 4 = 4

Example 3

Root 1 and leaf 3:

If 1 is bought, 3 gets discount:

  • cost(1)=4 profit=3
  • cost(3)=4 profit=7

Total cost = 8 ≤ 10

Total profit = 10

Example 4

Chain 1 → 2 → 3:

Buying all yields cascading discounts:

  • 1: 5 cost, profit 3
  • 2: 1 cost, profit 4
  • 3: 1 cost, profit 5

Total cost = 7, profit = 12

DP ensures chain propagation of discount state.

Complexity Analysis

Measure Complexity Explanation
Time O(n · budget²) Each node performs knapsack convolution over budget for children
Space O(n · budget) DP tables for two states per node and recursion stack

The quadratic budget factor arises from merging knapsack states at each node. Since each edge is processed once per state merge, total complexity remains bounded by n * budget².

Test Cases

# Provided examples
assert Solution().maxProfit(2, [1,2], [4,3], [[1,2]], 3) == 5  # basic chain with discount
assert Solution().maxProfit(2, [3,4], [5,8], [[1,2]], 4) == 4  # single best pick
assert Solution().maxProfit(3, [4,6,8], [7,9,11], [[1,2],[1,3]], 10) == 10  # root discount effect
assert Solution().maxProfit(3, [5,2,3], [8,5,6], [[1,2],[2,3]], 7) == 12  # chain full discount

# Edge cases
assert Solution().maxProfit(1, [10], [20], [], 5) == 10  # single node discounted choice trivial
assert Solution().maxProfit(3, [1,1,1], [2,2,2], [[1,2],[1,3]], 0) == 0  # zero budget
assert Solution().maxProfit(3, [10,10,10], [11,11,11], [[1,2],[1,3]], 1) == 0  # insufficient budget
assert Solution().maxProfit(4, [2,2,2,2], [5,5,5,5], [[1,2],[1,3],[3,4]], 6) >= 6  # mixed structure
Test Why
single node validates base case handling
zero budget ensures DP respects budget constraint
insufficient budget ensures no invalid selections
balanced tree validates multi-branch knapsack merge
chain structure validates discount propagation

Edge Cases

One important edge case is when the budget is zero. In this situation, the DP must correctly propagate that no purchases are possible, and all profit values remain zero. This tests whether initialization correctly prevents invalid transitions.

Another edge case occurs in a linear chain. Since discounts propagate strictly downward, the decision at the root affects every descendant. This stresses correctness of state propagation between parent and child in a single path.

A final edge case is a star-shaped tree where the root has many children. This stresses the knapsack merging logic, since multiple independent subtrees must be combined correctly under a shared budget constraint. Any mistake in convolution order or initialization would be exposed here.