atcoder/abc060D (feedback from editorial)

Author: xirc

atcoder/abc060/D/main.cpp

@@ -7,24 +7,36 @@ int N, W;
 vector<int> wss, vss;
 
 int solve() {
-    map<int,int> dp;
-    dp[0] = 0;
+    vector<deque<int>> v(4);
 
-    for (int i = 0; i < N; ++i) {
-        auto ndp = dp;
-        for (auto e : dp) {
-            auto w = e.first, v = e.second;
-            int nw = w + wss[i];
-            int nv = v + vss[i];
-            if (nw > W) continue;
-            ndp[nw] = max(ndp[nw], nv);
+    v[0].push_back(vss[0]);
+    for (int i = 1; i < N; ++i) {
+        int dw = wss[i] - wss[0];
+        assert(dw >= 0 && dw <= 3);
+        v[dw].push_back(vss[i]);
+    }
+
+    for (int i = 0; i < 4; ++i) {
+        sort(v[i].begin(), v[i].end());
+        reverse(v[i].begin(), v[i].end());
+        v[i].push_front(0);
+        for (int j = 1; j < v[i].size(); ++j) {
+            v[i][j] += v[i][j-1];
         }
-        dp = ndp;
     }
 
     int ans = 0;
-    for (auto e : dp) {
-        ans = max(ans, e.second);
+    for (int i = 0; i < v[0].size(); ++i) {
+        for (int j = 0; j < v[1].size(); ++j) {
+            for (int k = 0; k < v[2].size(); ++k) {
+                for (int l = 0; l < v[3].size(); ++l) {
+                    ll ww = ll(wss[0])*i + ll(wss[0]+1)*j + ll(wss[0]+2)*k + ll(wss[0]+3)*l;
+                    int vv = v[0][i] + v[1][j] + v[2][k] + v[3][l];
+                    if (ww > W) continue;
+                    ans = max(ans, vv);
+                }
+            }
+        }
     }
     return ans;
 }