Truly Optimal O(N) Time: Quick Select + Counting Sort
Intuition\n### Observation 1: How to quantify happiness :)\nIf we have selected k children, their total happiness is \n\n$\sum_{i = 0}^{k-1} max(0, \text{select
# Intuition\n### Observation 1: How to quantify happiness :)\nIf we have selected k children, their total happiness is \n\n$\\sum_{i = 0}^{k-1} max(0, \\text{selectedChildren}[i]- i)$\n\nWith this formula, we see that we can just select topK children in sorted order and arrive at our solution.\n\n$ans = \\sum_{i = 0}^{k-1} max(0, \\text{reverseSortedChildren}[i]- i)$\n\nSince we have to sort, this runs in the worst case $O(n\\log(n))$.\n\nHowever, with a few more observations we can see that we can do better.\n\n### Observation 2: When order doesn\'t matter (QuickSelect)\nIf all of the selected $k$ children have a $\\text{happiness} \\geq k$, their total happiness is \n\n$\\sum_{i = 0}^{k-1} (\\text{selectedChildren}[i]- i)$\n\nThis is because $(\\text{selectedChildren}[i] - i)$ will never be less than $1$ since $(i \\lt k \\leq \\text{selectedChildren}[i])$. Because of this, we can decouple:\n\n$\\sum_{i = 0}^{k-1} \\text{selectedChildren}[i] - \\sum_{i = 0}^{k-1} i$\n\nThis shows that the order we choose them doesn\'t matter.\n\nThus, if we take the $\\text{count\\_geq\\_k}$ of children with $\\text{happiness} \\geq k$, and see that $\\text{count\\_geq\\_k} \\geq k$, then we just take the topK happiest children in any order.\n\n$\\text{ans} = \\sum_{i = 0}^{k-1} (\\text{topKHappiest}[i] - i)$\n\nSince order doesn\'t matter, we can partition the $\\text{topKHappiest}$ using the $O(n)$ algorithm $\\text{quickselect}$. Technically, worst case performance is $O(n^2)$, but if we absolutely care about theoretical complexity we could use the $\\text{median of medians}$ algorithm to guarantee $O(n)$ performance. In the example below, I just choose a random pivot for simplicity.\n\n### Observation 3: When order does matter (CountingSort)\n\nSo what about the case where the $\\text{count\\_geq\\_k} \\lt k$?\n\nThis means all children with $\\text{happiness} \\geq k$ should automatically be included in our set and it doesn\'t matter which order we choose them.\n\nThe remaining children to be selected have $\\text{happiness} \\lt k$, which means the order we choose them does matter. However, since they are bounded by $k$, we can $\\text{counting sort}$ them in $O(k)$ which at worst is $O(n)$.\n\n# Complexity\n- Time complexity: $O(n)$\n\n- Space complexity: $O(n)$\n\n# Code\n```cpp\nclass Solution {\npublic:\n    long long maximumHappinessSum(vector<int>& happiness, int k) {\n        int geq_k_count = 0, n = happiness.size();\n        for (int x : happiness) {\n            geq_k_count += x >= k;\n        }\n\n        if (geq_k_count >= k) {\n            auto start = happiness.begin();\n            std::nth_element(start, start + n - k, happiness.end());\n            long long ans = 0;\n            for (int i = 0; i < k; ++i) {\n                ans += happiness[n - i - 1] - i;\n            }\n            return ans;\n        }\n\n        long long ans = 0, picked = 0;\n        std::vector<int> count(k);\n        for (int x : happiness) {\n            if (x < k) {\n                ++count[x];\n            } else {\n                ans += x - picked++;\n            }\n        }\n\n        for (int i = k - 1; i >= picked; --i) {\n            while (i >= picked && count[i]--) {\n                ans += i - picked++;\n            }\n        }\n        return ans;\n    }\n};\n```
['C++']