题目:
求n个数中前k个数的和
思路:
类似快排 减治法 使用随机基准数来优化复杂度
AC代码:
// #define NDEBUG
#include <bits/stdc++.h>
using namespace std;
namespace AIN
{
using db = double;
using ll = long long;
#define edl '\n'
#define str string
#define pll pair<ll, ll>
#define fir first
#define sec second
#define heap priority_queue
#define SPO(n) fixed << setprecision(n)
#define FOR(i, l, r) for (auto i = l; i <= (r); ++i)
#define ROF(i, r, l) for (auto i = r; i >= (l); --i)
#ifdef debugcmd
#define DBG(n) cout << "!!! " << #n << ": " << n << edl
#else
#define DBG(n) ;
#endif
template <typename T>
T KSM(T base, long long exp)
{
// assert(exp >= 0);
T res = 1;
while (exp)
{
if (exp & 1)
res *= base;
base *= base;
exp >>= 1;
}
return res;
}
// long long __lcm(long long lhs, long long rhs) { return lhs / __gcd(lhs, rhs) * rhs; }
mt19937_64 rnd(time(NULL));
long long RND(long long l, long long r) { return rnd() % (r - l + 1) + l; }
// constexpr db PI = acos(-1.0);
// constexpr db EPS = 1.0e-9;
constexpr long long LNF = 0x3f3f3f3f3f3f3f3fLL;
constexpr int INF = 0x3f3f3f3f;
// constexpr long long MOD = 998244353;
constexpr long long MOD = 1e9 + 7;
constexpr int MXN = 1e5 + 5;
}
using namespace AIN;
ll n, k;
ll ans;
ll a[MXN];
inline bool CMP(const ll lhs, const ll rhs)
{
return lhs <= rhs;
}
ll Sum(ll l, ll r)
{
ll res = 0;
FOR(i, l, r)
res += a[i];
return res;
}
void TopKSum(ll l, ll r, ll k)
{
if (l > r || k < 1 || k > (r - l + 1))
return;
swap(a[RND(l, r)], a[l]);
ll base = l;
ll pl = l;
ll pr = r;
while (pl < pr)
{
while (pl < pr && CMP(a[base], a[pr]))
--pr;
while (pl < pr && CMP(a[pl], a[base]))
++pl;
swap(a[pl], a[pr]);
}
swap(a[base], a[pl]);
ll cnt = pl - l + 1;
if (cnt <= k)
ans += Sum(l, pl);
TopKSum(pl + 1, r, k - cnt);
TopKSum(l, pl - 1, k);
return;
}
void Solve(void)
{
cin >> n >> k;
FOR(i, 1, n)
cin >> a[i];
TopKSum(1, n, k);
cout << ans << edl;
return;
}
int main(void)
{
std::ios::sync_with_stdio(false), cin.tie(nullptr), cout.tie(nullptr);
// #ifdef cincoutcmd
// freopen("cin.txt","r",stdin);
// freopen("cout.txt","w",stdout);
// #endif
Solve();
return 0;
}