POJ-3685 : Matrix

問題概要

N*N(N<5e4)の行列がある。A[i,j] = i*i+j*j+i*j+(i-j)*1e5であるとき、この行列の小さい方からM番目の値を求める問題。

解法

「小さい方からM番目がx」は「x以下の個数がM個以上となるような最小のx」と読み替えられるのでこれで二分探索する。で、x以下の個数を効率的に数え上げたい。iを[1..N]で動かし、各行について個数を求める。このとき、iを固定するとjの二次関数になっているのでunimordalとなり、ふたつの単調な列に分けることができる。それぞれで二分探索してやればよい。

acceptされたコード

#include <cstdio>
#include <algorithm>
using namespace std;

typedef long long int64;
const int64 INF = int64(1.05e18); 

int64 N, M;

void init() {
	scanf("%lld%lld", &N, &M);
}

int64 f(int64 y, int64 x) {
	return y*y + int64(1e5)*y + x*x - int64(1e5)*x + y*x;
}

bool check(int64 test) {
	int64 count = 0;
	for (int y = 1; y <= N; ++y) {
		int64 k = int64(1e5) - y;
		int64 mid = (k - 1) / 2;
		//[mid, N]で数える 増加列
		if (mid <= N) {
			mid = max<int64>(mid, 0);
			int64 low = mid - 1, high = N + 1;
			for (;high - low > 1;) {
				int64 m = (high + low) / 2;
				(f(y, m) > test ? high : low) = m;
			}
			count += high - mid;
		}
		//[1, mid)で数える 減少列
		if (1 < mid) {
			mid = min<int64>(mid, N + 1);
			int64 low = 0, high = mid;
			for (;high - low > 1;) {
				int64 m = (high + low) / 2;
				(f(y, m) > test ? low : high) = m;
			}
			count += mid - (low + 1);
		}
	}

	return count >= M;
}

int64 solve() {
	int64 low = -INF, high = INF;
	for (;high - low > 1;) {
		int64 mid = (high + low) / 2;
		(check(mid) ? high : low) = mid;
	}
	return high;
}

int main() {
	int T; scanf("%d", &T);
	for (int _ = 0; _ < T; ++_) {
		init();
		printf("%lld\n", solve());
	}
	return 0;
}