CodeForces - 1114B Yet Another Array Partitioning Task

题目链接

https://codeforces.com/problemset/problem/1114/B

解题思路

从大到小排序，找到最大的k * m个数；
答案就是这k * m个数的和；
遍历未排序的数组，每找到m个前k * m大的数就划成一段区间。

（woc？这么简单我没想到？？？？）

注意：
不可将排序后的第k * m个数当作下界去划分m个数；必须要选取前k * m个数作为划分数的集合。举个例子：
排序前：2 2 2 2 5 5 5 1
排序后：5 5 5 2 2 2 2 1
假设k = 4， m = 1，即分成4段，每段选取最大的数，排序后第第k * m个数为2，若以2为下界去将未排序的数组进行划分，则划分方式为2 2 2 2 5 5 5 1，此时答案为2+2+2+5 = 11，但是正确的划分方法为2 2 2 2 5 5 5 1 ，正确答案为2+5+5+5 = 17。
这个反例也是我实现大佬的思路时犯的错，前k * m个数是哪几个就只能用这几个数划分，不可以根据第k * m个数确定数的下界来划分！

代码

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 2e5+10;

struct node{
	int idx;
	LL num;
};
bool cmp(node x, node y) {
	return x.num>y.num?true:false;
}

node a[N], b[N];
LL ans = 0;
int tag[N];

int main()
{
	int n, m, k;
	scanf("%d%d%d", &n, &m, &k);
	for(int i = 1;i <= n;i ++) scanf("%lld", &a[i].num), a[i].idx = b[i].idx = i, b[i].num = a[i].num;
	
	sort(b+1, b+n+1, cmp);
	for(int i = 1;i <= k*m;i ++) ans += b[i].num, tag[b[i].idx] = 1;
	printf("%lld\n", ans);
	int cnt = 0, cc = 0;
	for(int i = 1;i <= n;i ++) {
		cnt += tag[a[i].idx];
		if(cnt == m) cnt = 0, printf("%d ", i), cc ++;
		if(cc == k-1) break;
	}

	return 0;
}