Find the Best Connected Strength for Every Tree Node
Company: Microsoft
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: medium
Interview Round: Take-home Project
## Find the Best Connected Strength for Every Tree Node
A tree contains `n` neurons. Each neuron is either strong or weak. Assign weight `+1` to a strong neuron and `-1` to a weak neuron.
For each neuron `v`, define its strength as the maximum total weight of any connected, nonempty set of tree vertices that contains `v`. Return the strength of every neuron.
### Function Signature
```python
def neuron_strengths(
n: int,
edges: list[tuple[int, int]],
strong_connectivity: list[int],
) -> list[int]:
...
```
### Input
- Vertices are numbered `1` through `n`.
- `edges` contains exactly `n - 1` undirected pairs and forms a tree.
- `strong_connectivity[i]` describes vertex `i + 1`: `1` is strong and `0` is weak.
### Output
Return an array of length `n`; element `i` is the maximum connected-set weight for vertex `i + 1`.
### Constraints
- `1 <= n <= 200_000`
- `strong_connectivity[i]` is `0` or `1`.
### Example
```text
n = 4
edges = [(1, 2), (1, 3), (1, 4)]
strong_connectivity = [0, 0, 1, 0]
Output: [0, -1, 1, -1]
```
For vertex `1`, choosing vertices `{1, 3}` gives weight `-1 + 1 = 0`. For vertex `2`, the best connected choice has total weight `-1`.
### Clarifications
- The chosen connected set may be different for each vertex.
- A set containing only `v` is always allowed, so every strength is well-defined even when all neurons are weak.
Quick Answer: For every node in a tree of strong and weak neurons weighted +1 and -1, find the best connected nonempty vertex set containing that node. A linear-time rerooting dynamic program combines positive subtree contributions and propagates useful value from the parent side.
For each vertex of a tree weighted +1 for strong and -1 for weak, return the maximum weight of a nonempty connected vertex set containing that vertex.
Constraints
- 1 <= n <= 200000
- edges forms a tree on vertices 1 through n
- strong_connectivity contains only 0 and 1
Examples
Input: {'n': 1, 'edges': [], 'strong_connectivity': [1]}
Expected Output: [1]
Explanation: A single strong neuron contributes one.
Input: {'n': 1, 'edges': [], 'strong_connectivity': [0]}
Expected Output: [-1]
Explanation: A single weak neuron must still be selected.
Hints
- Compute the best contribution from each rooted subtree.
- A second traversal can transfer the useful contribution from a node's parent side.