FastPrepTree Points

Tree Points

The D.E. Shaw Group logoThe D.E. Shaw Group● HardOA
Learn

Problem statement

Given a tree of g_nodes nodes rooted at 0, each node has a value given in array A. There is also a parameter, K.

An unweighted undirected tree is represented with g_nodes nodes numbered from 0 to g_nodes - 1 and g_edges = g_nodes - 1 edges where the ith edge connects the nodes numbered g_from[i] to g_to[i]. Each node numbered u is associated with a value represented by A[u].

Collect all of the values at the nodes using one of two methods:

  • Collect A[j] - K points.
  • Collect floor(A[j]/2) points. If this method is chosen, all values of node j in this node's subtree are changed to the floor(A[j]/2) as well.

You can only collect the value at a node if its parent's value has been collected before (except the root). You have to maximize the total points collected.

Note: It is possible that the accumulated value may be negative at some point.

Function

treePoints(g_nodes: int, g_from: int[], g_to: int[], A: int[], K: int) → long

Complete the function treePoints in the editor.

treePoints has the following parameters:

  1. int g_nodes: the number of nodes in the graph
  2. int g_from[g_nodes - 1]: one end of each edge
  3. int g_to[g_nodes - 1]: the other end each edge
  4. int A[g_nodes]: the values at the nodes
  5. int K: the cost to collect a full value

Returns long: the maximum sum of values that can be collected

Examples

Example 1

g_nodes = 4g_from = [0, 1, 2]g_to = [1, 2, 3]A = [10, 10, 3, 3]K = 5return = 11
Example 1 illustration

The first line shows the initial tree. Nodes are labeled as <node number>:<node value>. Darker nodes are ones that change in an operation.

For the first two values, reduce them by K = 5 and collect them. When the value at node 2 is collected as floor(3/2), the value at nodes j below it are changed using the same method, floor(A[j]/2). Hence, the value at node 3 is changed to 1. The value at node 3 is then collected using the second method as floor(1/2). The total collected is 5 + 5 + 1 + 0 = 11.

Constraints

  • 1 ≤ g_nodes ≤ 10^5
  • g_edges = g_nodes - 1
  • 0 ≤ g_from[i], g_to[i] < g_nodes
  • g_from[i] ≠ g_to[i]
  • 1 ≤ A[i] ≤ 10^9
  • 1 ≤ K ≤ 10^9

More The D.E. Shaw Group problems

See The D.E. Shaw Group hiring insights
public long treePoints(int g_nodes, int[] g_from, int[] g_to, int[] A, int K) {
  // write your code here
}
g_nodes4
g_from[0, 1, 2]
g_to[1, 2, 3]
A[10, 10, 3, 3]
K5
expected11
Checking account…