USACO 1.4 Arithmetic Progressions (堆排序 + 枚举)

#include <stdio.h>
#define DEBUG 1
#define TESTCASES 9
#define MAX 250

int lenOfProgressions;
int upperBound;
//双平方数的哈希
int isBisquare[2 * MAX * MAX + 1];
int arrayOfBisquares[2 * MAX * MAX + 1];
int numOfBisquares;
int heapSize;

void maxHeapify(int parent){
	int largest = parent;
	int left = parent << 1;
	if (left <= heapSize && arrayOfBisquares[left] > arrayOfBisquares[largest])
		largest = left;
	//"<<" 的优先级很低，注意加括号
	int right = (parent << 1) + 1;
	if (right <= heapSize && arrayOfBisquares[right] > arrayOfBisquares[largest])
		largest = right;
	if (largest != parent){
		int tempBisquare = arrayOfBisquares[largest];
		arrayOfBisquares[largest] = arrayOfBisquares[parent];
		arrayOfBisquares[parent] = tempBisquare;
		maxHeapify(largest);
	}
}

void buildMaxHeap(){
	int indexOfBisquare;
	for (indexOfBisquare = heapSize >> 1; indexOfBisquare >= 1; indexOfBisquare--)
		maxHeapify(indexOfBisquare);
}

void heapSortBisquares(){
	buildMaxHeap();
	int indexOfBisquare;
	for (indexOfBisquare = heapSize; indexOfBisquare > 1; indexOfBisquare--){
		int tempBisquare = arrayOfBisquares[indexOfBisquare];
		arrayOfBisquares[indexOfBisquare] = arrayOfBisquares[1];
		arrayOfBisquares[1] = tempBisquare;
		heapSize--;
		maxHeapify(1);
	}

}

int main(){
#if DEBUG
	int testCase;
	for (testCase = 1; testCase <= TESTCASES; testCase++){
		char inputFileName[20] = "inputx.txt";
		inputFileName[5] = '1' +  (testCase - 1);
		freopen(inputFileName, "r", stdin);
		printf("\n#%d\n", testCase);
#endif

	scanf("%d%d", &lenOfProgressions, &upperBound);

	int max = upperBound * upperBound * 2;	
	int num;
	for (num = 0; num <= max; num++)
		isBisquare[num] = 0;

	numOfBisquares = 0;
	int p, q;
	for (p = 0; p <= upperBound; p++)
		//因为p^2 + q^2跟q^2 + p^2是一样大的，所以q的枚举只要从p值开始就可以了
		for (q = p; q <= upperBound; q++){
			int bisquare = p * p + q * q;
			arrayOfBisquares[++numOfBisquares] = bisquare;
			isBisquare[bisquare] = 1;
		}
	
	heapSize = numOfBisquares;
	heapSortBisquares();
	
	int hasNoProgressions = 1;
	//因为a + (n - 1) * b = p ^ 2 + q ^ 2 = max,所以b < max / (lenOfProgressions - 1)
	int limitOfB = max / (lenOfProgressions - 1);
	int b;
	//因为输出要先按b排序所以先枚举b
	for (b = 1; b <= limitOfB; b++){
		int indexOfBisquare;
		//再按双平方数表枚举a
		for (indexOfBisquare = 1; arrayOfBisquares[indexOfBisquare] <= (max - (lenOfProgressions - 1) * b); indexOfBisquare++){
			int a = arrayOfBisquares[indexOfBisquare];
			int isProgressions = 1;
			int n;
			//从数列的最后一项开始往前枚举
			for (n = lenOfProgressions - 1; n > 0; n--)
				if (!isBisquare[a + n * b]){
					isProgressions = 0;
					break;
				}
			if (isProgressions){
				hasNoProgressions = 0;
				printf("%d %d\n", a, b);
			}
		}
	}
	if (hasNoProgressions)
		printf("NONE\n");

#if DEBUG
	}
#endif
	return 0;
}