CF:Codeforces Round 882 (Div. 2)A. The Man who became a God

Problem - A - Codeforces

AC代码：

#include<iostream>
#include<vector>
#include<algorithm>
#include<set>
#include<cmath>
#include<numeric>
using namespace std;
#define endl '\n'
#define ll long long
#define all(x) (x).begin(),(x).end()
#define rall(x) (x).rbegin(),(x).rend()
#define pii pair<int,int>
const int N = 2e5 + 100;
void func()
{
	int n, k;
	cin >> n >> k;
	vector<int>a(n);
	vector<pii>dif(n-1);
	vector<int>vis(n, 0);
	for (auto& x : a)cin >> x;
	int ans = 0;
	int zu = n;
	if (k == n)//一组一个人
	{
		//k==1==n||k==n>1
		ans = 0;
	}
	else//k<n
	{
		if (k == 1)//一个组且k!=n即n>1
		{
			for (int i = 0; i + 1 < n; i++)
			{
				ans += abs(a[i + 1] - a[i]);
			}
		}
		else//1<k<n
		{
			for (int i = 0; i  < n-1; i++)
			{
				dif[i].first=abs(a[i + 1] - a[i]);
				dif[i].second = i;
			}
			//dif[0]--->0-1,dif[1]--->1-2
			sort(dif.begin(), dif.end());
			for (int i = 0; i < n - 1; i++)
			{
				if (i == 0)
				{
					ans += dif[0].first;
					vis[dif[0].second] = 1;
					vis[dif[0].second + 1] = 1;
					zu--;
				}
				else
				{
					if (vis[dif[i].second] == 0 && vis[dif[i].second + 1] == 0)
					{
						ans += dif[i].first;
						vis[dif[i].second] = 1;
						vis[dif[i].second + 1] = 1;
						zu--;
					}
					else if (vis[dif[i].second] == 0 && vis[dif[i].second + 1] == 1)
					{
						ans += dif[i].first;
						vis[dif[i].second] = 1;
						zu--;
					}
					else if (vis[dif[i].second] == 1 && vis[dif[i].second + 1] == 0)
					{
						ans += dif[i].first;
						vis[dif[i].second+1] = 1;
						zu--;
					}
					else
					{
						ans += dif[i].first;
						zu--;
					}
				}
				if (zu == k)break;
			}
		}
	}
	cout << ans << endl;
}
int main()
{
	cin.tie(0), cout.tie(0)->sync_with_stdio(false);
	int t;
	cin >> t;
	while (t--)
	{
		func();
	}
	return 0;
}