Medium
Maximum Star Sum of a Graph — C++
Full explanation · Time O(n) · Space O(n)
// Time: O(n)
// Space: O(n)
// quick select
class Solution {
public:
int maxStarSum(vector<int>& vals, vector<vector<int>>& edges, int k) {
vector<vector<int>> adj(size(vals));
for (const auto& e : edges) {
if (vals[e[1]] > 0) {
adj[e[0]].emplace_back(e[1]);
}
if (vals[e[0]] > 0) {
adj[e[1]].emplace_back(e[0]);
}
}
int result = numeric_limits<int>::min();
for (int u = 0; u < size(vals); ++u) {
if (1 <= k && k <= size(adj[u])) {
nth_element(begin(adj[u]), begin(adj[u]) + k - 1, end(adj[u]),
[&](const auto& a, const auto& b) {
return vals[a] > vals[b];
});
}
result = max(result, vals[u] + accumulate(cbegin(adj[u]), cbegin(adj[u]) + min(k, static_cast<int>(size(adj[u]))), 0,
[&](const auto& total, const auto& x) {
return total + vals[x];
}));
}
return result;
}
};