TopK问题O（n）解法

题目：

求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;
}