Hard

Subtree Inversion Sum

LAB

Given tree edges and node values, return the maximum total sum after optionally inverting subtree signs with no two inverted roots adjacent.

EXAMPLES

Example 1

Input
{
  "values": [
    1,
    -2,
    3
  ],
  "edges": [
    [
      0,
      1
    ],
    [
      0,
      2
    ]
  ]
}

Output
6

FUNCTION SHAPE

values: intArrayedges: intMatrix→int

SOLUTION NOTE

Track distance since last inversion and parity (odd = currently inverted). Can only invert if dist >= k.

Reveal reference solution +

pythonREFERENCE

def subtreeInversionSum(self, edges: List[List[int]], nums: List[int], k: int) -> int:
    adjList = defaultdict(list)
    for u, v in edges:
        adjList[u].append(v)
        adjList[v].append(u)

    @cache
    def dp(root, par, dist, odd):
        no_invert = -nums[root] if odd else nums[root]
        invert = -nums[root] if not odd and dist >= k else nums[root]

        for nbr in adjList[root]:
            if nbr == par: continue
            no_invert += dp(nbr, root, dist + 1, odd)
            if dist >= k:
                invert += dp(nbr, root, 1, not odd)

        return max(no_invert, invert) if dist >= k else no_invert

    return dp(0, -1, float('inf'), False)

TimeO(n × k)

SpaceO(n × k)

Open on LeetCode ↗