Skip to content

Instantly share code, notes, and snippets.

@jedavidson
Last active January 24, 2024 06:44
Show Gist options
  • Save jedavidson/21a0679483dbdfc05396c0facbefca92 to your computer and use it in GitHub Desktop.
Save jedavidson/21a0679483dbdfc05396c0facbefca92 to your computer and use it in GitHub Desktop.
A mathematically satisfying explanation for a solution to the "Stick Lengths" problem from the CSES problem set

This is the problem in question.

Let $\ell_i$ be the length of the $i^{\textsf{th}}$ stick. Without loss of generality, we will assume that $$\ell_1 \leq \ell_2 \leq \cdots \leq \ell_n.$$ (The eventual solution has no need for the lengths to be sorted, so this is just for convenience in the correctness proof.)

Making all sticks have length $\ell$ has associated cost $$c(\ell) = \sum_{i = 1}^{n} |\ell_i - \ell|,$$ and the problem asks us for the minimal cost over all lengths. We may as well ask instead for any optimal length $\ell^*$, since we can just compute $c(\ell^*)$ to get the minimal cost if need be.

If you have a bit of statistics knowledge, this immediately gives the game away: this is a sum of absolute deviations in $L_1$ norm from some given data points, and this is precisely minimised about a median of those points, so we just need to compute a median. Let's pretend we don't know that though and derive it for ourself.

The algorithm

Suppose we only have two sticks (i.e. $n = 2$). For any $\ell \in [\ell_1, \ell_2]$, we have $$c(\ell) = (\ell - \ell_1) + (\ell_2 - \ell) = \ell_2 - \ell_1,$$ and $c(\ell) > \ell_2 - \ell_1$ for $\ell < \ell_1$ or $\ell > \ell_2$. Thus, any length $\ell^* \in [\ell_1, \ell_2]$ is optimal. This should make intuitive sense as well, because any more length we need to add to the shorter stick is an equal amount less length we have to take from the second stick. Moreover, this equilibrium is disturbed whenever we try and make the sticks both longer or both shorter than they started out.

We now consider the general case, though we will first treat the case where $n$ is even. Pairing up the $i^{\textsf{th}}$ shortest stick with the $i^{\textsf{th}}$ largest stick (i.e. the $(n + 1 - i)^{\textsf{th}}$ stick) gives us $n / 2$ pairs of sticks. For the $i^{\textsf{th}}$ pair, let $$c_i(\ell) = |\ell_{i} - \ell| + |\ell_{n + 1 - i} - \ell|$$ be the associated cost of making those sticks sticks have length $\ell$. The previous analysis of the two stick case shows that $c_i$ is minimised in the interval $[\ell_i, \ell_{n + 1 - i}]$. But since these intervals form a descending chain, i.e. $$[\ell_1, \ell_n] \supseteq [\ell_2, \ell_{n - 1}] \supseteq \cdots \supseteq [\ell_{n/2}, \ell_{n / 2 + 1}],$$ picking any $\ell^* \in [\ell_{n/2}, \ell_{n / 2 + 1}]$ will in fact minimise all $c_i$, so $$c(\ell) = \sum_{i = 1}^{n/2} c_i(\ell) \geq \sum_{i = 1}^{n} c_i(\ell^*) = c(\ell^*).$$ When $n$ is odd, the above pairing process gives us $\lfloor n / 2 \rfloor$ pairs with the stick of unique median length left over, so in this case $$c(\ell) = |\ell_{(n + 1) / 2} - \ell| + \sum_{i = 1}^{\lfloor n / 2 \rfloor} c_i(\ell).$$ But it is clear that $\ell_{(n + 1) / 2} \in [\ell_{\lfloor n / 2 \rfloor}, \ell_{\lfloor n / 2 \rfloor + 2}]$, which is the innermost minimising interval with respect to the sum above, so $$c(\ell) \geq |\ell_{(n + 1) / 2} - \ell_{(n + 1) / 2}| + \sum_{i = 1}^{\lfloor n / 2 \rfloor} c_i(\ell_{(n + 1) / 2}) = c(\ell_{(n + 1) / 2}).$$

We can now make the following completely general statement: if we have $n$ sticks, the optimal length to make them is

  • $\ell_{(n + 1) / 2}$ if $n$ is odd, or
  • $\ell_{n/2}$ if $n$ is even.

These are the median and left-of-median lengths respectively, which agrees with the statistics intuition.

The implementation

Seeing as programming languages typically index starting 0, and since integer division typically rounds down by default, the length we are actually interested in finding in either case is $\ell_{n / 2}$. One easy approach is to sort the lengths and pick out the value at index $n / 2$, but a better (on average) solution is to use quickselect, which can find the median in expected $O(n)$ time.

Some languages have quickselect implemented already, like C++. Hence, our solution is

int minimal_cost(std::vector<int>& lengths) {
    // Find the median
    auto m = lengths.size() / 2;
    std::nth_element(lengths.begin(), std::next(lengths.begin(), m), lengths.end());
    int median = lengths[m];
    
    // Compute the minimal cost
    int cost = 0;
    for (int l : lengths) {
        cost += std::abs(l - median);
    }

    return cost;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment