Problem analysis: Discipline (stage III, XXXIII OI)

Abridged problem statement

We have \(n\) heaps of stones of sizes \(a_1, a_2, \ldots, a_n\). Every day we can choose at most one heap and take one stone from it. However, there is an additional constraint: on day \(k\) (for \(k =1,2,\dots\)) we can take a stone from heap \(i\) only if in the binary representation of \(k\), its \(i\)-th least significant bit is set. Our goal is to find the minimum number of days needed to take all the stones, and the result should be output modulo \(10^9+7\).

Subtask 1: \(n\leq 4,\ a_i\leq 5\) — backtracking

The simplest approach is recursive backtracking, which every day considers taking a stone from each of the heaps allowed on that day, or not taking any stone. We can add the following optimization here: if we are currently at day \(d\), we still have a total of \(s\) stones left, and \(d+s\) is greater than the best result found so far, we can immediately stop exploring this branch, because we will not find a better solution in it than the one we currently have.

In the recursive calls, we must remember to first consider all options for taking stones, and only at the end consider not taking any, to avoid infinite recursion.

Subtask 2: \(n\leq 6,\ a_i\leq 10\) — minor optimization

We can easily speed up the above approach as follows. If on a given day we have the option to take a stone from some heap, it is always no worse than not taking any stone. In this case, we can immediately skip the option of not taking a stone, which significantly reduces the number of recursive calls.

Subtask 3: \(n\leq 8,\ a_i\leq 30\) — Matchings

In this problem, we are looking for the smallest number \(d\) such that it is possible to match one of the days \(1, \ldots, d\) to each stone. To check if a given value \(d\) is sufficient, we can create a bipartite graph where the vertices on one side represent the stones (there are \(\sum a_i\) of them), and the vertices on the other side represent the days from 1 to \(d\) (there are \(d\) of them). An edge between a stone and a day exists if the given stone can be taken on the given day.

Then, using an algorithm for finding the maximum matching in a bipartite graph, we can check if there exists a matching where each stone is assigned to some day.

Therefore, we can check if a given value \(d\) is sufficient for consecutive values of \(d\) starting from \(\sum a_i\) upwards, until we find the first value for which such a matching exists.

Alternatively, we can also binary search over \(d\), however, given the constraints of this subtask, linear checking is fast enough. The idea of binary searching over \(d\) will be very useful in subsequent subtasks, though.

Subtask 4: \(n\leq 8\) — flow network

In this subtask, both the number of stones and the number of days are too large for standard matching checking. Note, however, that if \(d\) is a sufficient number of days to move all the stones, then any number of days greater than \(d\) is also sufficient. Thus, we can binary search the answer. This reduces the problem to checking if a given number of days is sufficient.

To do this, we can group all stones of a given heap into a single vertex, and also group the days by their remainder modulo \(2^n\), since days with the same remainder have edges to the same heaps.

We build a flow network on these vertices, where the edges have the following capacities: \(a_i\) units flow from the source to the heaps, the capacity between heaps and days is infinite (if a stone can be taken from a given heap on a given day), and the number of units flowing from the days to the sink is equal to the number of days among \(1,2,\dots,d\) with the given remainder modulo \(2^n\). On such a network, we can use Dinic's or Edmonds-Karp algorithm to find the maximum flow and check if it is exactly equal to \(\sum a_i\). If so, the given value \(d\) is sufficient; otherwise, it is not.

Subtask 5: \(n\leq 16\) — weighted Hall's theorem

For such large inputs, the previously created flow network is too large for any flow algorithm. We will therefore use a different approach, based on the weighted Hall's theorem.

Recall the classic Hall's Marriage Theorem: in a bipartite graph with vertices \(L\) on the left side and \(R\) on the right, a matching covering all vertices of \(L\) exists if and only if for every subset \(S \subseteq L\), \(|S| \leq |N(S)|\) holds, where \(N(S)\) is the set of neighbors of the vertices in \(S\).

The weighted version of Hall's theorem generalizes this condition to the case where vertices have non-unit weights. Let each vertex \(v \in L\) have a weight \(w_L(v)\), and each vertex \(u \in R\) have a weight \(w_R(u)\). We are looking for a flow in which exactly \(w_L(v)\) units flow out of each \(v \in L\), and at most \(w_R(u)\) units flow into each \(u \in R\). The weighted Hall's theorem states that such a flow exists if and only if for every subset \(S \subseteq L\) we have \[\sum_{v \in S} w_L(v) \leq \sum_{u \in N(S)} w_R(u).\]

In our problem, the vertices on the left are the heaps with weights \(a_i\), and on the right are the day groups (grouped by their remainder modulo \(2^n\)) with weights equal to the number of days of a given type among \(1, 2, \ldots, d\). For a fixed value of \(d\) from the binary search, we must check whether for every subset of heaps \(S\), \(\sum_{i \in S} a_i \leq f(S)\) holds, where \(f(S)\) denotes the total weight of the day groups adjacent to the heaps in \(S\).

The number of days of each type can be easily calculated based on \(d\). Then, to determine \(f(S)\) for all subsets \(S\), we perform SOS (sum-over-subsets) dynamic programming, which allows us to compute these values in \(O(n \cdot 2^n)\) complexity.

Combining this with binary search over the answer, we get a solution with \(O(n \cdot 2^n \cdot \log \text{ans})\) complexity, which is sufficient for this subtask.

Subtask 6: \(n\leq 20\) — direct calculation in \(O(n \cdot 2^n)\)

In this subtask, we can get rid of the binary search and directly calculate for each subset \(S\) the minimum number of days required so that the days "serving" this subset are enough to remove all its corresponding stones.

The algorithm for a single subset \(S\) is as follows. We represent the subset \(S\) as a bitmask of length \(n\), where bit \(i\) is set if and only if \(i \in S\). Let \(m = \max(S)\) be the largest element of the subset (meaning we can only take stones from heap \(m\) on days where the \(m\)-th bit is set). We index the bits from 1 to \(n\), where 1 is the least significant bit.

Let's look at a cycle of days of length \(2^m\). Let's consider how many days in this cycle have at least one bit belonging to \(S\) set. Notice that during such a cycle, every possible bitmask of length \(m\) appears exactly once on the \(m\) least significant bits of the day number. Let \(s = |S|\). During this cycle, we have exactly \((2^s-1)\cdot 2^{m-s}\) days in which at least one bit from \(S\) is set. The factor \(2^s - 1\) is the number of non-empty configurations of the \(s\) bits belonging to \(S\) (at least one bit must be set, which subtracts one possibility from all \(2^s\) possibilities), and \(2^{m-s}\) is the number of configurations of the remaining \(m-s\) bits (those outside \(S\), but below \(m\)), which can be anything.

Having this value, we subtract as many full cycles as possible from the sum of stones \(\sum_{i \in S} a_i\), increasing the number of required days by the length of the cycle (\(2^m\)) for each subtraction. We are left with a remainder of stones and must consider an incomplete cycle. Let's consider the first half of the cycle, i.e., days \(1, \ldots, 2^{m-1}-1\). Among these days there are exactly \((2^{s-1}-1)\cdot 2^{m-s}\) days with a set bit from \(S\), because in this half of the cycle the \(m\)-th bit is not set, so we only have \(s-1\) bits from \(S \setminus \{m\}\) available, and the remaining \(m-s\) bits can be arbitrary.

If this number of days is smaller than the number of stones remaining to be served, we know we must use all these days and then proceed to the second half of the cycle, i.e., days \(2^{m-1}, \ldots, 2^m-1\). On each of these days the \(m\)-th bit is set, so we can use each of them to take a stone. Therefore, in this case, we add \(2^{m-1}\) (the length of the first half of the cycle) to the answer, plus as many days from the second half of the cycle as there are stones left to serve after using the days from the first half.

Otherwise, if the number of remaining stones is no greater than \((2^{s-1}-1)\cdot 2^{m-s}\), we know we will no longer use the \(m\)-th bit, so we remove it from \(S\) (i.e., we consider \(S' = S \setminus \{m\}\)) and repeat the procedure recursively for \(S'\) with a new \(m' = \max(S')\).

Calculating the value for a single subset takes \(O(n)\) time, so for all subsets we get an \(O(n \cdot 2^n)\) complexity. The final answer in the problem is the maximum over all subsets.

Subtask 7: \(a_i\leq 10^5\) — lemma for large \(n\)

It is easy to see that for a sufficiently large \(n\), only \(a_n\) matters. We can formalize this with the following lemma:

Lemma: If \(2^{n-2} \geq \sum_{i=1}^{n-1} a_i\), and \(a_n \leq 2^{n-1}\), then the answer to the problem is \(2^{n-1} + a_n - 1\).

Proof: Note that the result is at least \(2^{n-1}\), because this is the number of the first day on which we can take a stone from heap \(n\). Let's consider days from 1 to \(2^{n-1} - 1\) and heaps from \(1\) to \(n-1\). For each of these heaps, among the considered days, there are exactly \(2^{n-2}\) days on which we can take stones from it. Using the weighted version of Hall's theorem again, we can see that the condition mentioned in this theorem is satisfied for every subset, because we assign each stone from these heaps at least as many days as there are stones in total. Thus, there is some assignment of all stones from heaps \(1\) to \(n-1\) to days from \(1\) to \(2^{n-1} - 1\). After assigning these stones, we are left with at most \(a_n\) stones from heap \(n\) to handle, and we have days from \(2^{n-1}\) onwards available, so we can handle them all on days from \(2^{n-1}\) to \(2^{n-1} + a_n - 1\), each of which is good for heap \(n\), because \(a_n \leq 2^{n-1}\).

With the constraints in this subtask, using the lemma, we can print the result according to the formula for \(n \geq 24\), and use the solution from subtask 6 for smaller \(n\).

Model solution — digit DP

In the main version of the problem, the lemma allows us to limit \(n\) to about 57, which is still too large for the previous approach of iterating over all subsets. We therefore need to look closer at what happens in the Hall's theorem approach.

Let \(d\) be the value from the binary search. For a subset of heaps \(S\), let \(f(S)\) denote the number of days between 1 and \(d\) (inclusive) that have edges to \(S\). We want to check if for any subset \(S\), \(\sum_{i \in S} a_i > f(S)\) holds, which is equivalent to checking if \(\sum_{i \in S} a_i + (d + 1 - f(S)) > d + 1\). The expression \(d + 1 - f(S)\) is the number of days among the days from 0 to \(d\) that have no edges in the considered graph (there are no edges to day number 0).

If we find a subset \(S\) that maximizes the expression \(\sum_{i \in S} a_i + (d + 1 - f(S))\), we will simply be able to check if this value exceeds \(d+1\). If so, \(d\) is too small, and if not, \(d\) is sufficient.

Let \(g(S) = d + 1 - f(S)\) denote the number of days among \(0, 1, \ldots, d\), that do not share any common bit with the mask \(S\). Note that for full cycles \([0, 2^n-1], [2^n, 2 \cdot 2^n - 1], \ldots\), the number of days disjoint from the mask \(S\) is always \(2^{n - |S|}\). Let \(d + 1 = c \cdot 2^n + r\), where \(c\) is the number of full cycles, and \(r \in [0, 2^n)\) is the remainder. Then \[g(S) = c \cdot 2^{n - |S|} + h(S, r),\] where \(h(S, r)\) is the number of days in the interval \([0, r)\) disjoint from \(S\).

To calculate \(h(S, r)\), we will consider the binary representations of the numbers. For a number \(x\), we denote by \(x_i\) its \(i\)-th bit, where we number the bits from 1 (least significant) to \(n\) (most significant). Similarly, \(S_i\) denotes whether position \(i\) belongs to mask \(S\) and likewise for \(r_i\).

We iterate over all possible positions \(p\), which are the highest position where both \(S\) and \(r\) have a one. Fixing such a position \(p\) imposes constraints on the mask \(S\):

At position \(p\): We must have \(S_p = 1\) (by definition of \(p\)).
Above \(p\) (positions \(i > p\)): If \(r_i = 1\), then we must have \(S_i = 0\) (otherwise \(p\) would not be the highest common one). If \(r_i = 0\), we can freely choose \(S_i \in \{0, 1\}\).
Below \(p\) (positions \(i < p\)): There are no restrictions and we can freely choose \(S_i \in \{0, 1\}\).

Let's consider what numbers \(x < r\) disjoint from \(S\) look like. Every such number \(x\) must have \(x_p = 0\) (because \(S_p = 1\)), which means that the longest common prefix of \(x\) and \(r\) ends at a position greater than \(p\). At positions below \(p\), the number \(x\) can have any bits, as long as they are disjoint from \(S\).

For positions \(p, p+1, \dots, n\) we define the following dynamic programming: \[\text{DP}(i, j) = \text{maximum value of } \sum_{k \in S} a_k + |\{x < r : x \cap S = \emptyset,\, x \text{ shares a common prefix with } r \text{ down to position } i+1\}|\] among all \(S \subseteq \{1, \ldots, i\}\), where \(|S| = j\) and \(p\in S\).

First, let's consider how to calculate \(\text{DP}(p, j)\) for \(j \in [1, p]\). Naturally, \(S\) must contain position \(p\), where there is a one, hence we consider \(j \geq 1\). So we have \(j-1\) ones to distribute among positions \(1, \ldots, p-1\). The number of days \(x\) that share a common prefix with \(r\) at least down to position \(p+1\) and are disjoint from \(S\), is exactly \(2^{p-j}\), because at positions where we do not choose ones in \(S\), we can have any bits, and every such binary sequence corresponds to a value less than \(r\) (because it has a common prefix down to position \(p+1\) and a 0 at position \(p\), where \(r\) has a 1). Therefore, \[ \text{DP}(p, j) = 2^{p-j} + a_p + \text{sum of } j-1 \text{ largest } a_k \text{ for } 1 \leq k < p. \]

Next, calculating \(\text{DP}(i, j)\) for \(i > p\), we have two possibilities. If \(r_i = 1\), we know we must have \(S_i = 0\). Then no value \(a_i\) is added to the sum, and the number of days disjoint from \(S\), less than \(r\), and sharing a common prefix with \(r\) at least down to position \(i+1\), increases relative to that counted in \(\text{DP}(i-1, j)\). This is because we now allow days \(x\) where \(x_i\) can be 0, whereas previously we did not, as \(x_i\) had to be the same as \(r_i\). This number increases by exactly \(2^{i-1-j}\), because at positions \(1, \dots, i-1\), where we do not choose ones in \(S\), we can have any bits. Thus, we get the formula \[ \text{DP}(i, j) = \text{DP}(i-1, j) + 2^{i-1-j}. \]

If, however, \(r_i = 0\), we can have \(S_i = 1\) or \(S_i = 0\). In the first case, we add \(a_i\) to the sum, and the number of days \(x\) disjoint from \(S\), less than \(r\), and sharing a common prefix with \(r\) at least down to position \(i+1\), does not change (new days must have \(x_i = 0\), and we also have \(r_i = 0\)). In the second case, we do not add \(a_i\) to the sum, but the number of days also does not change, because new days must have \(x_i = 0\), as they must be the same as \(r\) at positions greater than \(i\) and must be less than \(r\) (so they cannot have \(x_i = 1\)). We thus obtain the formula \[ \text{DP}(i, j) = \max(\text{DP}(i-1, j-1) + a_i, \text{DP}(i-1, j)). \]

Finally, it is enough to look at \(\text{DP}(n, j)\) for \(j \in [1, n]\) and find the maximum among the values \(\text{DP}(n, j) + c \cdot 2^{n-j}\), because our value \(\text{DP}(n, j)\) corresponds to the calculated \(h\) restricted only to masks in which exactly \(j\) bits are set, and the highest common bit of \(S\) and \(r\) is at position \(p\).

By repeating this algorithm for all positions \(p\) from \(1\) to \(n\), where \(r\) has a set bit, as well as for the special case where we consider masks \(S\) disjoint from \(r\) (we solve it with a similar dynamic programming approach that starts from an artificial position \(p = 0\)), we can find the maximum value \(\sum_{i \in S} a_i + g(S)\) in \(O(n^3)\) complexity for a single \(d\). Combined with binary search over \(d\), this yields a solution with \(O(n^3 \cdot \log(\text{ans}))\) complexity, where \(\text{ans}\) is the answer to the problem.