Hi there! So, I was trying to solve 1954D - Colored Balls. When I first looked at it, I thought "omg how do i solve this fast? iz impossibruh!!1!". And then I looked at the additional constraint and be like "bruh it's just knapsack". Now, after thinking about the idea to solve this problem and then looked at the editorial (turns out my idea was more complex lmao even though it's the same concept), I was wondering, "ain't no way knapsack can solve the problem but without additional constraint. r-right?". Well, I got good news. Users, I present to you, the solution of 1954D - Colored Balls but the total number of balls can actually exceed $$$5000$$$!

Now, If you want to understand what I'm expanding on, you should check out the editorial for the problem first, as I am too lazy to write everything here and I also have limited time before I have to go back to play Hypixel Skyblock (not sponsored btw) and then sleep. Now that's out of the way, here's how to solve it:

To recap what the editorial has said:

1. Denote the sum of a subset as $$$s$$$. Each subset's value is $$$\big\lceil \frac{s}{2} \big\rceil$$$ (we call this type 1 subset), unless the maximum integer inside that subset is larger than that. In that case, the subset's value is the maximum integer itself (we call this type 2 subset).
2. Assume all subsets are type 1 subset. Then, we can compute the number of subsets with sum $$$j$$$ using knapsack. Then, compute the sum of values over all subsets.
3. Ofc sometimes not every subset is a type 1 subset. And so we have to compensate the missing values for all type 2 subsets. So, for each $$$i$$$ such that $$$1 \leq i \leq n$$$, we compute the number of subsets such that $$$a_i > \big\lceil \frac{s}{2} \big\rceil$$$. Then, for each of that subset, we compensate it by $$$a_i - \big\lceil \frac{s}{2} \big\rceil$$$ (this can easily be done using the knapsack array we already created from 2nd point).

Now, it seems like the bottleneck for this is the 2nd point, where we have to do knapsack in $$$O(nS)$$$ where $$$S$$$ is the total number of balls. Since, if we do it exactly like the editorial, it'll TLE (cuz $$$S$$$ could reach like $$$25 \cdot 10^6$$$ lmao). But, notice that $$$a_i \leq 5000$$$. So, why not use knapsack to just compute number of subsets with sum from $$$0$$$ to sum $$$\max(a_i)$$$ in $$$O(n \cdot a_i)$$$, and then use that knapsack to support 3rd point? "uhm, but KitasCraft, how are we suppose to count the sum of values over all subsets from the 2nd point then???" Well, we can use a really interesting trick here. Note that from this point on, all of the text are talking about modifying the 2nd point, so you still have to do the 3rd point normally afterward lol.

Now, let's modify the 2nd point a little bit. Assume that every "type 1 subset" has the value of $$$s$$$ instead of $$$\big\lceil \frac{s}{2} \big\rceil$$$. Now, the problem for the 2nd point has been reduced to counting sum over all subset sums. If you already know the "Contribution Sums" technique, then congrats. If you haven't, read this (although the blog applied it on a tree problem).

By using the "Contribution Sums" technique, we'll notice that each value of $$$a_i$$$ appears in exactly $$$2^{n-1}$$$ subsets. And so, we can count the sum over all subset sums in $$$O(n)$$$ using this technique. And, if we divide the sum by $$$2$$$, we can count the sum of values over all subsets in $$$O(n)$$$! Well, that is if the value is $$$\frac{s}{2}$$$ instead of $$$\big\lceil \frac{s}{2} \big\rceil$$$. So, how do we get past this?

Notice that $$$\big\lceil \frac{s}{2} \big\rceil = \frac{s}{2}$$$ if $$$s$$$ is even and $$$\big\lceil \frac{s+1}{2} \big\rceil = \frac{s}{2}$$$ if $$$s$$$ is odd. Thus, we only need to count the sum over all subset sums, add it by the number of subsets with odd sum, then divide all of those by $$$2$$$. This way, we can properly count the sum of values over all subsets in $$$O(n)$$$.

Observation. The number of subsets with odd sum is either $$$0$$$ (if there are no odd $$$a_i$$$) or $$$2^{n-1}$$$ (if otherwise)

Proof. Assume that we have a set $$$T$$$ of size $$$k$$$ with only even numbers. Then, it is very obvious that the number of subsets with odd sum and even sum is $$$0$$$ and $$$2^k$$$ respectively. If we add an odd number into $$$T$$$ (now the size of $$$T$$$ is $$$k+1$$$), then it is guaranteed that every subset that has this number will result in odd sum. And so, the number of subsets with odd sum and even sum is $$$2^k$$$ and $$$2^k$$$ respectively. If we add another odd number into $$$T$$$ (now the size of $$$T$$$ is $$$k+2$$$), then a subset $$$T'$$$ that contains this number actually has the reversed parity of the subset of $$$T'$$$ that does not contain this number. And so, if we apply this into all of the other subsets of $$$T$$$ that contains the new number, then the number of subsets with odd sum is $$$2^k$$$ (the old subset with odd sum A.K.A. the subset with odd sum and without the new number) added by $$$2^k$$$ (the new subset with odd sum A.K.A. the subset that used to have even sum, but then the new number was added, thus making this subset has odd sum), which is $$$2^{k+1}$$$. Same goes for the number of subsets with even sum.

And so, using those observation, we can actually optimize the knapsack from $$$O(nS)$$$ to $$$O(n \cdot a_i)$$$ and we can also optimize the 2nd point from $$$O(nS)$$$ to $$$O(n)$$$. Yayyy! oh btw here's the code 256647759

P.S. My nature language is not english. So, there might be mistakes inside this blog. Feel free to point some mistakes in this blog and also to give feedbacks :3 (help i've written this for like an hour and it's like 10 p.m. now lmao)