TCO Algorithm Final Editorial

The finalists for this round are tourist, Errichto, Um_nik, ACRush, qwerty787788, ksun48, Petr, _aid.

This round, the easy is a problem that requires a greedy observation and is simple to implement afterwards. The medium is a twist on a spanning tree problem on a 2d plane, and requires strong geometric intuition to convince yourself the solution is correct. The last problem is a tricky dp problem that can get very messy if the edge cases are not handled in a sane manner.

BalancingTrees (lg5293):

In this problem, we are given a rooted tree where there are weights at each node. We call a node “weight balanced” if all of its child subtrees have the same total weight. A tree is called weight balanced if every node is weight balanced. We are allowed to change the weights at each node, and the cost of such a change is the absolute difference between the old and new weights. We want to find the minimum cost needed to make the tree weight balanced.

First, we claim we only ever need to change the value at the leaves of the tree. To see this, suppose we wanted to change the weight of a non leaf node x. We can instead propagate this change down to the children of x divided evenly, and this is still a solution with at most the same cost which keeps all children with the same weight (we can keep propagating these changes down until we hit the leaves).

Let’s consider a general function f(x, W) which is the cost needed in node x’s subtree to make its subtree have weight . Due to the observation above, we can compute this greedily by pushing all our changes to the leaves. For a non leaf node i, we can compute f(i, W) as sum of , over the children x of node i. For a leaf node i, we can compute f(i, W) as |W – wi|.

For general T, as we compute this dp downwards, we can see that the second argument is a linear function in T. For each node, we can compute this linear function as we go down. The cost at each leaf node is of the form |ai · T + bi – wi|. Thus, we can reduce this problem to one where we are given a bunch of lines, and we want to find a value t such that the sum of the absolute values of the lines at point t is minimized.

There are a few ways to do this. We can use ternary search since this is the sum of convex functions. Another approach is to notice that we only need to check values t where a line hits zero, so there are only O(n) values to check.

Overall time complexity is O(n2) or O(n · log MAXi), which is enough to pass this problem.

There is also a faster solution in O(n) using observations above. Looking for the optimal t can be viewed as a weighted median problem: you have some people on the line who want to meet in the same spot, person i currently stands at , and moving them one unit of distance costs ai. This can be done by using quick select to find any point at which the slope changes from nonpositive to nonnegative.

A sample solution using ternary search is shown below:

import java.util.ArrayList;
import java.util.List;
public class BalancingTrees {

    int n;
    int[] p,w;
    List<Integer>[] child;
    double[] want;
    public double minCost(int[] _p, int[] _w) {
        p = _p; w = _w;
        n = w.length;
        child = new List[n];
        for (int i = 0; i < n; i++) {
            child[i] = new ArrayList<>();
        for (int i = 0; i < n-1; i++) {
        want = new double[n];

        double lo = -(1L<<40), hi = (1L<<40);
        for (int iter = 0; iter < 200; iter++) {
            double f1 = (2*lo+hi)/3.0, f2 = (lo+2*hi)/3.0;
            double x1 = getValue(f1), x2 = getValue(f2);
            if (x1 < x2) {
                hi = f2;
            } else {
                lo = f1;

        return getValue((lo+hi)/2.0);

    double getValue(double rootval) {
        want[0] = rootval;
        double ret = 0;
        for (int i = 0; i < n; i++) {
            want[i] -= w[i];
            int nc = child[i].size();
            for (int nxt : child[i]) {
                want[nxt] = want[i] / (double)nc;
            if (nc == 0) {
                ret += Math.abs(want[i]);
        return ret;

BuildTheRoads (monsoon):

We have n cities on a two-dimensional plane (each with two coordinates); no three cities lie on one line. We need to connect them with roads in such a way that there is exactly one path between every pair of cities and the length of the longest such path is minimal (length of road is Euclidean).

The crucial observation in this problem is that there is an optimal solution in which there is no simple path with more than three edges.

Proof: suppose that there is such a path A1 – B1 – C1 – … – C2 – B2 – A2 (possibly with C1 = C2 ) that A1 and A2 are leaves and there’s no v such that d(v, Bi) > d(Ai, Bi).

Wlog we assume that d(A1, B1) + d(B1, C1) ≤ d(A2, B2) + d(B2, C2).

Then we replace all edges d(v, B1) for v ≠ C1 with edge d(v, C1), thus making vertex B1 a leaf.

This won’t increase diameter, since for every path in new tree we can find the same (or longer) path in the old tree.

The only problematic paths are from A1; let v be connected to C1 (either B1 or outside path); in new graph

So we produced a tree with no bigger diameter, but with one more leaf. Since we cannot increase number of leaves in infinity, finally we end up with a tree without a path of more than three edges.

Trees without a path of more than three edges can be described as follows: we have some edge (A, B) and all other vertices are connected either to A or B (this also includes star graphs).

We are now ready for an algorithm.

We iterate over O(n) possibilities of choosing A and B.

Then we iterate over O(n2) possibilities of choosing vertex X connected to A which maximizes d(A, X).

We claim that for all vertices v, we connect v to A if d(v, a) ≤ d(A, X) and to B otherwise.

To prove that this is an optimal solution, let X = R2 and ,let d(A, X) = r2, and consider a point R1 such at distance r1 < r2 from A, and this was the second longest node directly connected to A. Everything with d(v, A) ≤ r1 can be connected directly to A without increasing the diameter.

Now, suppose there was a point P such that r2 ≤ d(A, P) ≤ r1. Then, we claim that in an optimal solution we can connect P directly to A as well.

We claim that connecting P to A will not make the diameter worse than connecting P to B. For instance, consider the maximum length path from node R2. It can only be better to connect P to A directly by triangle inequality. Similarly, consider the maximum length path from a leaf of B. Since R2 is already the maximum length leaf from A, it does not matter which node P is directly connected to. Thus, it is optimal to only consider fix the maximum and take everything inside that circle to one point.

That leads to O(n4) algorithm. It can be improved: fix A and B.

For every other vertex v calculate (d(A,v),d(B,v)) and sort these pairs non-increasingly over first coordinate.

Then iterate over them: if we consider i-th pair we assume that all later pairs go to A (they have no bigger d(A,v)) and all the previous pairs go to B (so we keep running maximum over second coordinate).

We need to carefully calculate second maximums. This gives us an O(n3 log n solution.

We can get rid of a log factor since we can observe that we don’t need O(n2 sorts, but only O(n) (we only sort over first coordinate which is the same for every O(n) sorts).

So we can preprocess sorting which results in O(n3 algorithm, but this was not required.

import java.util.*;

public class BuildTheRoads {

public class DistPair implements Comparable<DistPair> {
  int from_i, from_j;
  DistPair(int from_i, int from_j) {
    this.from_i = from_i;
    this.from_j = from_j;
  public int compareTo(DistPair other) {
    return other.from_i - this.from_i;  // descending

public int sq(int x) { return x*x; }

public double minimalCost(int[] x, int[] y) {
  int n = x.length;
  if (n == 2) { return Math.sqrt(sq(x[0] - x[1]) + sq(y[0] - y[1])); }

  double best = 1_000_000_000;

  // The main lemma: there is an optimal tree in which every path has length at most 3.
  // Such tree has one edge (i,j) and every other vertex is connected either to i or j.

  for (int i=0; i<n; ++i) for (int j=i+1; j<n; ++j) {

    // We choose every candidate for (i,j) and sort other vertices descending according to pair
    // (squared distance from i, squared distance from j).

    int dist = sq(x[i] - x[j]) + sq(y[i] - y[j]);

    ArrayList<DistPair> dists = new ArrayList<DistPair>();
    for (int k=0; k<n; ++k) if (k != i && k != j) {
      dists.add(new DistPair(sq(x[i] - x[k]) + sq(y[i] - y[k]), sq(x[j] - x[k]) + sq(y[j] - y[k])));

    // Maximum and second-maximum squared distances from i (respectively j) to other vertices
    // connected to i (j).

    double first_i = 0, second_i = 0;
    double first_j = 0, second_j = 0;

    // We choose every candidate for first_i. All vertices with non-greater distance from i than
    // first_i are connected to i. Other vertices are connected to j.

    for (int k=0; k < dists.size(); ++k) {
      first_i = dists.get(k).from_i;
      second_i = k+1 < dists.size() ? dists.get(k+1).from_i : 0;

      double cost = Math.sqrt(first_i) + Math.sqrt(dist) + Math.sqrt(first_j);
      cost = Math.max(cost, Math.sqrt(first_i) + Math.sqrt(second_i));
      cost = Math.max(cost, Math.sqrt(first_j) + Math.sqrt(second_j));

      best = Math.min(best, cost);

      double val = dists.get(k).from_j;
      if (val > first_j) { second_j = first_j; first_j = val; }
      else if (val > second_j) { second_j = val; }

  return best;


Worms (misof):

First, we define a worm in a 2D grid to be a sequence of cells that starts at one cell, and only goes up or right. We would like to partition the 2D grid into some worms, where each cell appears in exactly one worm. We are also given a partial solution, and we want to count the number of ways to complete the partial solution. The partial grid is given in a special form: if we arrange all the cells in row-major order then a prefix of this order will be provided.

The main observation of is problem is when drawing worms onto an empty board, each main diagonal is independent.


More precisely, each cell of the diagonal must belong to a different worm, and we can make an independent set of choices where for each cell we decide whether the worm in this cell continues up, right, or terminates in this cell. We just need to make sure that two worms don’t collide.

We can count the number of choices for each diagonal using a simple dynamic programming.
The state is described by the length of the diagonal and by two booleans: whether we can
choose “up” for the leftmost cell and whether we can choose “right” for the rightmost cell. In other words, this is a 1d dynamic programming problem, where we want to choose a string of length n of characters “U” and “R” and “S” (for “up”, “right”, and “stop” respectively), where we cannot have the substring “RU” in our string, and we may have extra restrictions “the string cannot start in U” and “the string cannot end in R”.

Thus, when counting the number of worm decompositions for an empty grid, all we need is the
above DP to determine the number of choices for each diagonal and then to multiply those numbers.

If the grid has a prefix, there are three more cases to consider:



The worm in the leftmost cell of the diagonal shown above cannot go “up”. In general, you cannot go up if the cell above you is not the leftmost known cell of that worm.



In the above grid we are missing a part of the “c” worm. We need to check for this and fill in the rest of the worm:


If we have a partially-filled row and case 2 did not apply:


the “c”-worm can go arbitrarily far to the right. One way to handle this
is to try all possibilities for the length of the “c”-worm and for each of them use
the solution for the empty grid. As the solution sets are disjoint, we can simply
sum them up.

import java.util.*;

public class Worms {

    boolean isWorm(char[][] board, char ch) {
        int R = board.length, C = board[0].length;
        int[][] where = new int[R+C-1][2];
        String fingerprint = "";
        for (int diag=0; diag<R+C-1; ++diag) {
            boolean found = false;
            for (int r=0; r<R; ++r) {
                int c = r+diag-R+1;
                if (!(0 <= c && c < C)) continue;
                if (board[r][c] != ch) continue;
                if (found) return false; // two on the same diagonal
                found = true;
                where[diag][0] = r;
                where[diag][1] = c;
            fingerprint += found ? 'y' : 'n';
        if (!fingerprint.matches("n*y*n*")) return false;
        for (int diag=0; diag<R+C-2; ++diag) {
            if (fingerprint.charAt(diag) == 'n' || fingerprint.charAt(diag+1) == 'n') continue;
            boolean adjacent = false;
            if (where[diag+1][0] == where[diag][0] && where[diag+1][1] == where[diag][1]+1) adjacent = true;
            if (where[diag+1][0] == where[diag][0]-1 && where[diag+1][1] == where[diag][1]) adjacent = true;
            if (!adjacent) return false;
        return true;

    int[] lastLetter(char[][] board) {
        int R = board.length, C = board[0].length;
        int rl = 0, cl = 0;
        if (board[0][0] != '?') {
            while (true) {
                if (rl == R-1 && cl == C-1) break;
                int nr = rl, nc = cl+1;
                if (nc == C) { ++nr; nc = 0; }
                if (board[nr][nc] == '?') break;
                rl = nr; cl = nc;
        return new int[] {rl,cl};

    boolean fillInBrokenWorm(char[][] board) {
        int[] ll = lastLetter(board);
        int rl = ll[0], cl = ll[1];
        char last = board[rl][cl];
        if (last == '?') return true;
        if (isWorm(board,last)) return true;
        while (cl < board[rl].length-1) {
            board[rl][cl] = last;
            if (isWorm(board,last)) return true;
        return false;

    public int countThisBoard(char[][] board, int[][][] dp) {
        int R = board.length, C = board[0].length;
        long answer = 1;
        for (int diag=0; diag<R+C-1; ++diag) {
            int n=0, fl=0, fr=0;
            for (int r=0; r<R; ++r) {
                int c = r+diag-R+1;
                if (!(0 <= c && c < C)) continue;
                if (board[r][c] != '?') continue;
                if (n == 1) {
                    if (c == 0 && r > 0) fl = 1;
                    if (c > 0 && r > 0 && board[r-1][c] == '?') fl = 1;
                    if (c > 0 && r > 0 && board[r-1][c] != '?' && board[r-1][c-1] != board[r-1][c]) fl = 1;
                if (c == C-1) fr = 0; else fr = 1;
            answer *= dp[n][fl][fr];
            answer %= 1000000007;
        return (int)answer;

    public int count(String[] printed) {

        int R = printed.length, C = printed[0].length();
        char[][] board = new char[R][C];
        for (int r=0; r<R; ++r) for (int c=0; c<C; ++c) board[r][c] = printed[r].charAt(c);

        // fill in the broken worm, if any

        // precompute the dp for a single diagonal
        int[][][] dp = new int[51][2][2];
        for (int n=0; n<=50; ++n) for (int fl=0; fl<2; ++fl) for (int fr=0; fr<2; ++fr) {
            if (n == 0) { dp[n][fl][fr] = 1; continue; }
            dp[n][fl][fr] = 0;

            // we can always stop the leftmost worm, which allows its neighbor to go up
            dp[n][fl][fr] += dp[n-1][1][fr];

            // if we are allowed, we can continue the leftmost worm upwards, same effect
            if (fl == 1) dp[n][fl][fr] += dp[n-1][1][fr];
            dp[n][fl][fr] %= 1000000007;

            // if we are allowed (which we always are for n > 1), we can continue the leftmost worm to the right, which blocks neighbor from going up
            if (n > 1 || (n == 1 && fr == 1)) dp[n][fl][fr] += dp[n-1][0][fr];
            dp[n][fl][fr] %= 1000000007;

        // if last printed letter can continue to the right, try all possibilities how far
        int[] ll = lastLetter(board);
        int lr = ll[0], lc = ll[1];

        if (board[lr][lc] == '?') return countThisBoard(board,dp);
        if (lr > 0 && board[lr-1][lc] == board[lr][lc]) return countThisBoard(board,dp);

        long answer = countThisBoard(board,dp);
        while (true) {
            if (lc == C) break;
            board[lr][lc] = board[lr][lc-1];
            answer += countThisBoard(board,dp);
        answer %= 1000000007; 
        return (int)answer;