Update maximum-number-of-groups-getting-fresh-donuts.cpp

Author: kamyu104

C++/maximum-number-of-groups-getting-fresh-donuts.cpp

@@ -17,30 +17,45 @@ class Solution {
             count[i] -= pair_count;
             count[size(count) - i] -= pair_count;
         }
-        int max_mask = accumulate(cbegin(count), cend(count), 1,
-                                  [](int total, int c) {
-                                      return total * (c + 1);
+        unordered_map<int, int> new_count;
+        for (int i = 0; i < size(count); ++i) {
+            if (count[i]) {
+                new_count[i] = count[i];
+            }
+        }
+        int max_mask = accumulate(cbegin(new_count), cend(new_count), 1,
+                                  [](int total, pair<int, int> kvp) {
+                                      return total * (kvp.second + 1);
                                   });
         vector<int> lookup(max_mask);
-        return result + memoization(count, max_mask - 1, 0, &lookup);
+        return result + memoization(batchSize, new_count, max_mask - 1, 0, &lookup);
     }
 
 private:
-    int memoization(const vector<int>& count, int mask, int remain, vector<int> *lookup) {
+    int memoization(int batchSize, const unordered_map<int, int>& count, int mask, int remain, vector<int> *lookup) {
         if (!(*lookup)[mask]) {
-            int curr = mask, basis = 1, i = 0;
-            for (; i < remain; basis *= (count[i] + 1), curr /= (count[i] + 1), ++i);
-            // mask: a0 + a1 * (c0 + 1)  + a2 * (c0 + 1) * (c1 + 1) + ... + a(b-1) * (c0 + 1) * (c1 + 1) * ... * (c(b-2) + 1)
-            int a_remain = curr % (count[remain] + 1);
+            int a_remain = 0;
+            int curr = mask, basis = 1;
+            if (count.count(remain)) {
+                for (const auto& [i, c] : count) {
+                    if (i == remain) {
+                        break;
+                    }
+                    basis *= (c + 1), curr /= (c + 1);
+                }
+                // mask: a0 + a1 * (c0 + 1)  + a2 * (c0 + 1) * (c1 + 1) + ... + a(b-1) * (c0 + 1) * (c1 + 1) * ... * (c(b-2) + 1)
+                a_remain = curr % (count.at(remain) + 1);
+            }
             int result = 0;
             if (a_remain) {  // greedily use remain
-                result = max(result, (remain == 0) + memoization(count, mask - basis, 0, lookup));
+                result = max(result, (remain == 0) + memoization(batchSize, count, mask - basis, 0, lookup));
             } else {
-                for (int curr = mask, basis = 1, i = 0; i < size(count); basis *= (count[i] + 1), curr /= (count[i] + 1), ++i) {
-                    if (curr % (count[i] + 1) == 0) {
-                        continue;
+                int curr = mask, basis = 1;
+                for (const auto& [i, c] : count) {
+                    if (curr % (c + 1)) {
+                        result = max(result, (remain == 0) + memoization(batchSize, count, mask - basis, ((remain - i) + batchSize) % batchSize, lookup));
                     }
-                    result = max(result, (remain == 0) + memoization(count, mask - basis, ((remain - i) + size(count)) % size(count), lookup));
+                    basis *= (c + 1), curr /= (c + 1);
                 }
             }
             (*lookup)[mask] = result;
@@ -67,21 +82,28 @@ class Solution2 {
             count[i] -= pair_count;
             count[size(count) - i] -= pair_count;
         }
-        int max_mask = accumulate(cbegin(count), cend(count), 1,
-                                  [](int total, int c) {
-                                      return total * (c + 1);
+        unordered_map<int, int> new_count;
+        for (int i = 0; i < size(count); ++i) {
+            if (count[i]) {
+                new_count[i] = count[i];
+            }
+        }
+        int max_mask = accumulate(cbegin(new_count), cend(new_count), 1,
+                                  [](int total, pair<int, int> kvp) {
+                                      return total * (kvp.second + 1);
                                   });
         vector<int> dp(max_mask);
         for (int mask = 0; mask < size(dp); ++mask) {
             int remain = 0;
-            for (int curr = mask, basis = 1, i = 0; i < size(count);
-                 basis *= (count[i] + 1), curr /= (count[i] + 1), ++i) {  // decode mask
+            int curr = mask, basis = 1;
+                for (const auto& [i, c] : new_count) {
                 // mask: a0 + a1 * (c0 + 1)  + a2 * (c0 + 1) * (c1 + 1) + ... + a(b-1) * (c0 + 1) * (c1 + 1) * ... * (c(b-2) + 1)
                 int ai = curr % (count[i] + 1);
                 if (ai > 0) {
                     dp[mask] = max(dp[mask], dp[mask - basis]);
                 }
-                remain = (remain + ai * i) % size(count);
+                remain = (remain + ai * i) % batchSize;
+                basis *= (c + 1), curr /= (c + 1);
             }
             if (mask != size(dp) - 1 && remain == 0) {
                 ++dp[mask];