43 lines
1.1 KiB
C++
43 lines
1.1 KiB
C++
#include <vector>
|
|
#include <cmath>
|
|
#include <algorithm>
|
|
#include <iostream>
|
|
|
|
void print(std::vector<int>& gifts) {
|
|
for (auto gift : gifts)
|
|
std::cout << gift << ' ';
|
|
std::cout << "\n---\n";
|
|
}
|
|
|
|
namespace {
|
|
class Solution {
|
|
public:
|
|
long long pickGifts(std::vector<int>& gifts, int k) {
|
|
std::sort(gifts.begin(), gifts.end());
|
|
|
|
for (int i = 0; i < k; ++i) {
|
|
auto value = std::sqrtl(gifts[gifts.size() - 1]);
|
|
if (value == 1)
|
|
continue;
|
|
auto it = std::upper_bound(gifts.begin(), gifts.end(), value);
|
|
auto idx = it - gifts.begin();
|
|
std::shift_right(it, gifts.end(), 1);
|
|
gifts[idx] = value;
|
|
}
|
|
long long sum = 0;
|
|
for (auto gift : gifts) {
|
|
sum += gift;
|
|
}
|
|
return sum;
|
|
}
|
|
};
|
|
}
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
TEST(P2558, 1) {
|
|
std::vector<int> gifts = {1,1,1,1};
|
|
int k = 4;
|
|
Solution s;
|
|
auto res = s.pickGifts(gifts, k);
|
|
} |