Notation

It is not hard to observe that since for each village there is a unique way down the river from the village to Bytetown, we can treat the rivers and villages as a tree with the root in Bytetown. Nodes of the tree correspond to the villages (for convenience we will refer to Bytetown as a village too), and node ~v~ is the parent of node ~u~ when ~v~ is the first village downriver from ~u~.

Let ~r~ denote the root of the tree, i.e. ~r~ corresponds to Bytetown. By ~depth(u)~ we will denote the number of edges on a unique path from ~u~ to ~r~. Clearly, the values of ~depth(u)~ can be computed for all villages ~u~ in linear time. The number of children of node ~u~ will be denoted by ~\deg(u)~, and the number of trees cut near village ~u~ will be denoted by ~trees(u)~.

Dynamic Programming

We can apply dynamic programming to solve the task. Let ~A_{v,t,l}~ denote the minimal cost of transportation of the trees cut in the subtree rooted in ~v~, assuming that ~t~ additional sawmills can be built in the subtree, and the trees not processed in ~v~ can be processed in the village of depth ~l~ (on the way from ~v~ to Bytetown). We compute values of ~A_{v,t,l}~ for each village ~v~, and such numbers ~t,l~, that ~0 \le t \le k~ and ~0 \le l < depth(v)~. Clearly, when the tree rooted in ~v~ has at most ~t~ nodes, then ~A_{v,t,l} = 0~, as we can simply place a sawmill in every village. We can use the following formula:

$$\displaystyle A_{v,t,l} = \begin{cases} 0 & \text{when the tree rooted in }v\text{ has at most }t\text{ nodes} \\ \min(A'_{v,t,l}, A''_{v,t}) & \text{otherwise} \end{cases} \tag{1}$$

where ~A'_{v,t,l}~ is the cost of transportation when there is no sawmill in ~v~, and ~A''_{v,t}~ is the cost of transportation when there is one. These costs depend on the distribution of sawmills between subtrees rooted in children of ~v~. Let ~d = \deg(v)~ and let ~v_1, v_2, \dots, v_d~ be the children of ~v~. Then:

$$\displaystyle A'_{v,t,l} = trees(v) \times (depth(v)-l) + \min_{t_1+\dots+t_d=t} \sum_{i=1}^d A_{v_i,t_i,l} \tag{2}$$

$$\displaystyle A''_{v,t} = \min_{t_1+\dots+t_d=t-1} \sum_{i=1}^d A_{v_i,t_i,depth(v)} \tag{3}$$

Dynamic Programming One More Time

Let us have a closer look at the recurrences (2) and (3). We do not need to consider every partition of ~t~ into a sum of ~\deg(v)~ terms, to compute the values of ~A'~ and ~A''~. It would be time-consuming in case of trees containing vertices with many children. Once again, we can use dynamic programming. Let ~B^{i,s}_{v,l}~ denote the cost of transporting the trees cut in subtrees rooted in ~v_1, v_2, \dots, v_i~ provided that ~s~ additional sawmills can be built in these subtrees and the trees not processed in these subtrees are processed in a village of depth ~l~. We can make use of the following recurrence:

$$\displaystyle \begin{align*} B^{0,s}_{v,l} &= 0 \\ B^{i,s}_{v,l} &= \min_{0 \le j \le s} (B^{i-1,s-j}_{v,l}+A_{v_i,j,l})\text{ for each }s = 1, \dots, k \end{align*} \tag{4}$$

We define ~C^{i,s}_v~ analogously, but this time we assume that the trees not processed in the subtrees are processed in ~v~. Then

$$\displaystyle \begin{align*} C^{0,s}_v &= 0 \\ C^{i,s}_v &= \min_{0 \le j \le s} (C^{i-1,s-j}_{v,l}+A_{v_i,j,depth(v)})\text{ for each }s = 1, \dots, k \end{align*} \tag{5}$$

Obviously, ~A'_{v,t,l} = B^{\deg(v),t}_{v,l}~ and ~A''_{v,t} = C^{\deg(v),t-1}_v~. In order to compute all values of ~A'_{v,t,l}~ (respectively ~A''_{v,t}~), for some ~l \in \{0, \dots, depth(v)\}~, we compute ~B^{i,s}_{v,l}~ (respectively ~C^{i,s}_v~) for each ~i = 0, \dots, \deg(v)~ and ~s = 0, \dots, k~. It follows that for each pair ~v,l~ it takes ~\mathcal O(k^2 (\deg(v)+1))~ time to compute values ~B^{i,s}_{v,l}~ and ~C^{i,s}_v~, giving total time:

$$\displaystyle \mathcal O\left(n \sum_v k^2 (deg(v)+1)\right) = \mathcal O(k^2 n^2)$$

since ~\sum_v \deg(v)~ is equal to the number of edges in the tree, i.e. ~n-1~. After computing all values of ~A_{v,t,l}~ the program returns ~A'(r,k,0)~ as the final answer.

Too Much Dynamic Programming

Using dynamic programming twice would not be required if all the vertices in the tree had a small number of children. Fortunately, we can easily construct a binary tree of villages that has the same minimal cost of transportation as the original one. In order to do this, we connect the first child of each vertex to the parent as its left child, create an additional village and connect it as a right child to the parent. Then we connect the other children of the parent vertex in the original tree one by one, each time connecting a child as a left child of the additional village created during connecting the previous child, creating a new additional village and connecting it as a right child of the previous one. Additional villages produce no wood and rivers connecting them to parents are of length ~0~, so they do not affect the total cost of transportation. Building sawmills in additional villages does not lower the minimal cost of transportation, because every sawmill built in an additional village can be moved to the next non-additional village on the way to Bytetown without raising (and even possibly lowering) the total cost.

There are ~t+1~ partitions of ~t~ into a sum of ~2~ non-negative terms. For each pair ~v,l~ we can therefore compute ~A_{v,t,l}~ for all ~t~ in ~\mathcal O(k^2)~ time. There are twice as many vertices in the binary tree as in the original one. The total computation time is:

$$\displaystyle \mathcal O\left(n \sum_v k^2\right) = \mathcal O(k^2 n^2)$$

One River

It is worth noting, that in the special case, where instead of a tree of rivers there is just one river with multiple villages, the problem is significantly simplified and can be solved more efficiently. In such a case we can use simple dynamic programming. Let us number the villages from ~1~ to ~n~ upstream. For each village ~v~, and for each number of sawmills ~q~ ~(0 \le q \le k)~ let us denote by ~T[v,q]~ the optimal cost of locating ~q~ sawmills in the part of the river downstream from ~v~. The value of ~T[v,q]~ can be computed in ~\mathcal O(v)~ time: we check all ~v~ possibilities of placing the upstream-most sawmill using stored values of ~T[1,q-1], \dots, T[v-1,q-1]~. Such an algorithm works in ~\mathcal O(kn^2)~ time.

A special case of such a problem — for just two sawmills — appeared at CEOI 2004, but the upper bound on ~n~ was much bigger: ~n \le 20\,000~. This required at least an ~\mathcal O(n \log n)~ algorithm, which did not use dynamic programming, but was based on the divide and conquer technique instead. There exists also a solution running in linear time, iterating over all possible locations of the uppermost sawmill and finding the optimal location of the other sawmill in amortized constant time.