leetcode

Leetcode Solutions (0x6A73)
git clone https://git.pngu.org/leetcode
Log | Files | Refs | Submodules | README

max_profit_trade_stocks.py (3449B)


# 3562. Maximum Profit from Trading Stocks with Discounts

"""
you are given an integer 'n' representing the number of employees in a
company. each employee is assign a unique id from 1 to 'n', and employee 1 is
the ceo. you are given two integers arrays 'present' and 'future', each of
length 'n', where 'present[i]' represents the current price at which the i'th
employee can buy  stock today. 'future[i]' represents the expected price at
which the i'th employee can sell the stock tomorrow. the company's hierarchy
is represented by a 2d integer array 'hierarchy', where 'hierarchy[i] = [u_i,
v_i]' means that employee 'u_i' is the direct boss of employee can buy their
stock at half the original price. return the maximum profit that can be
achieve without exceeding the budget.
"""


class Solution(object):
    def maxProfit(self, n, present, future, hierarchy, budget):
        """
        :type n: int
        :type present: List[int]
        :type future: List[int]
        :type hierarchy: List[List[int]]
        :type budget: int
        :rtype: int
        """
        tree = [[] for _ in range(n)]
        for u, v in hierarchy:
            tree[u - 1].append(v - 1)
        dp = [[[0] * (budget + 1) for _ in range(2)] for _ in range(n)]
        self.dfs(0, present, future, tree, dp, budget)
        return max(dp[0][0])

    def dfs(self, u, present, future, tree, dp, budget):
        children = tree[u]
        child_dps = []
        for v in children:
            self.dfs(v, present, future, tree, dp, budget)
            child_dps.append((dp[v][0], dp[v][1]))
        for parentBought in range(2):
            price = present[u] // 2 if parentBought else present[u]
            profit = future[u] - price
            base = [0] * (budget + 1)
            for c0, _ in child_dps:
                next_base = [0] * (budget + 1)
                for b in range(budget + 1):
                    if base[b] == 0 and b != 0:
                        continue
                    for k in range(budget - b + 1):
                        next_base[b + k] = max(next_base[b + k], base[b] + c0[k])
                base = next_base
            curr = base[:]
            if price <= budget:
                base = [0] * (budget + 1)
                for _, c1 in child_dps:
                    next_base = [0] * (budget + 1)
                    for b in range(budget + 1):
                        if base[b] == 0 and b != 0:
                            continue
                        for k in range(budget - b + 1):
                            next_base[b + k] = max(next_base[b + k], base[b] + c1[k])
                    base = next_base
                for b in range(price, budget + 1):
                    curr[b] = max(curr[b], base[b - price] + profit)
            dp[u][parentBought] = curr


if __name__ == "__main__":
    obj = Solution()
    print(
        obj.maxProfit(n=2, present=[1, 2], future=[4, 3], hierarchy=[[1, 2]], budget=3)
    )
    print(
        obj.maxProfit(n=2, present=[3, 4], future=[5, 8], hierarchy=[[1, 2]], budget=4)
    )
    print(
        obj.maxProfit(
            n=3,
            present=[4, 6, 8],
            future=[7, 9, 11],
            hierarchy=[[1, 2], [1, 3]],
            budget=10,
        )
    )
    print(
        obj.maxProfit(
            n=3,
            present=[5, 2, 3],
            future=[8, 5, 6],
            hierarchy=[[1, 2], [2, 3]],
            budget=7,
        )
    )