chore: sync automated leetcode submissions with Runtime - 84 ms (98.67%), Memory - 107.9 MB (67.44%)

Author: bufferwise

3351-maximize-happiness-of-selected-children/solution.cpp

@@ -1,10 +1,26 @@
+#include <vector>
+#include <algorithm>
+
+static const int fast_io = []() {
+    std::ios_base::sync_with_stdio(false);
+    std::cin.tie(NULL);
+    return 0;
+}();
+
 class Solution {
 public:
-    long long maximumHappinessSum(vector<int>& a, int k) {
-        long long res = 0, i;
-        sort(a.begin(), a.end());
-        reverse(a.begin(), a.end());
-        for (i = 0; i < k; ++i) res += max(0LL, a[i] - i);
-        return res;
+    long long maximumHappinessSum(std::vector<int>& happiness, int k) {
+        std::nth_element(happiness.begin(), happiness.begin() + k, happiness.end(), std::greater<int>());
+        std::sort(happiness.begin(), happiness.begin() + k, std::greater<int>());
+
+        long long totalHappiness = 0;
+        for (int i = 0; i < k; ++i) {
+            int currentHappiness = happiness[i] - i;
+            if (currentHappiness <= 0) break;
+            
+            totalHappiness += currentHappiness;
+        }
+        
+        return totalHappiness;
     }
 };