4815. 【NOIP2016提高A组五校联考4】ksum

题解

题目描述

给出一个长度为 nnn 的数组 aaa.

求 aaa 中所有的子段和中最小的 kkk 个.

Sample Input

复制代码

3 4
1 3 4

Sample Output

复制代码

8 7 4 4

数据范围

n,k≤105,0<ai≤109n,k\le 10^5,0<a_i\leq10^9n,k≤105,0<ai≤109

样例说明

现在列表计算子段和:

子段	和
{1}\{1\}{1}	111
{3}\{3\}{3}	333
{4}\{4\}{4}	444
{1,3}\{1,3\}{1,3}	444
{3,4}\{3,4\}{3,4}	777
{1,3,4}\{1,3,4\}{1,3,4}	888

数组是 {1,3,4,4,7,8}\{1, 3, 4, 4, 7, 8\}{1,3,4,4,7,8} 倒序就是 {8,7,4,4,3,1}\{8, 7, 4, 4, 3, 1\}{8,7,4,4,3,1}.

求前 k=4k=4k=4 个, 就是 {8,7,4,4}\{8, 7, 4, 4\}{8,7,4,4}.

思路

这种求前 kkk 小值的问题一般都是使用优先队列解决的归并问题.

因为 k≤105k\le10^5k≤105, 所以状态的转移一定是 O(1)O(1)O(1) 的.

我们可以计算出 aaa 的前缀和数组 preipre_iprei. 那么区间 [l,r][l,r][l,r] 的和就是 prer−prel−1pre_r-pre_{l-1}prer−prel−1.

因为 ai>0a_i>0ai>0 所以这个前缀和数组一定是递增的.

题意转化

那么现在题目就变成了: 给你一个递增的数组 preipre_iprei, 求 prei−prej−1(i>j)pre_i-pre_{j-1}(i>j)prei−prej−1(i>j) 的前 kkk 大的值.

那么这个就是一个归并的模板题目, 优先队列按照值排序即可, 总的时间复杂度为 O(klog⁡k)O(k\log k)O(klogk).

代码

cpp 复制代码

#include <stdio.h>

#define MAXN 100010
#define ri register int
#define FOR(i, a, b) for (ri i = (a); i < (b); i++)
#define REP(i, a, b) for (ri i = (a); i <= (b); i++)
#define For(i, a) FOR(i, 0, a)
#define Rep(i, a) REP(i, 1, a)

typedef long long LL;

int n, k;
LL PreSum[MAXN];
int arr[MAXN];

struct Node_t {
	int j, i;
	LL val;
};
int CompareNode(struct Node_t x, struct Node_t y) {return x.val > y.val;}
typedef struct Node_t Node, *PNode;
void swap(PNode x, PNode y) {Node t = *x; *x = *y; *y = t;}

// 堆
Node heap[MAXN << 2];		// 小心数组开小去世
int heap_size = 0;

// 向上调整
void HeapUp(int pos) {
	while (pos > 1 && !CompareNode(heap[pos >> 1], heap[pos])) {
		swap(&heap[pos >> 1], &heap[pos]);
		pos >>= 1;
	}
}
// 向下调整
void HeapDown(int pos) {
	while ((pos << 1) <= heap_size) {
		int t = pos << 1;
		if ((pos << 1 | 1) <= heap_size && !CompareNode(heap[t], heap[pos << 1 | 1])) t = pos << 1 | 1;
		if (CompareNode(heap[pos], heap[t])) break;
		swap(&heap[pos], &heap[t]), pos = t;
	}
}
// 堆顶
Node HeapTop() {return heap[1];}
int IsEmpty() {return heap_size == 0;}
void HeapPush(Node p) {heap[++heap_size] = p; HeapUp(heap_size);}
void HeapPop() {swap(&heap[1], &heap[heap_size]); heap_size --; HeapDown(1);}
void print() {
	while (!IsEmpty()) printf("%lld\n", HeapTop().val), HeapPop();
}

int main()
{
	freopen("ksum.in", "r", stdin);
	freopen("ksum.out", "w", stdout);
	
	scanf("%d%d", &n, &k);
	Rep(i, n) {
		scanf("%d", &arr[i]);
		PreSum[i] = PreSum[i - 1] + arr[i];
	}
	
	// 现在问题就成了: 求 pre[j] - pre[i - 1] 的前 k 大值, 归并即可
	Rep(i, n) HeapPush((Node) {i, 1, PreSum[i]});
	
	
	int print_val = 0;
	while (!IsEmpty()) {
		if (print_val >= k) break;
		print_val++;
		printf("%lld ", HeapTop().val);
		
		Node TopVal = HeapTop(); HeapPop();
		
		if (TopVal.i + 1 <= TopVal.j) HeapPush((Node) {TopVal.j, TopVal.i + 1, PreSum[TopVal.j] - PreSum[TopVal.i]});
	}
	
	return 0;
}