[백준/C++] 9426 - 중앙값 측정

정승우·2021년 2월 23일

[백준/C++] BOJ 공부

목록 보기

16/25

문제링크: https://www.acmicpc.net/problem/9426

문제

기상학에서 주요 사용하는 대표값은 중앙값이다. (중앙값의 정의는 힌트에 나와있다)

상근이는 1초에 한 번씩 온도를 재는 기계를 가지고 있고, 이 기계에 들어갈 소프트웨어를 작성하려고 한다. 기계에는 작은 디지털 디스플레이가 하나 달려있다. 매 초마다 디스플레이는 지난 K초동안 측정한 온도의 중앙값을 화면에 보여준다.

상근이는 소프트웨어를 기계에 올리기 전에 컴퓨터에서 테스트해보려고 한다.

총 N초 동안 측정한 온도가 주어졌을 때, 디스플레이에 표시된 중앙값의 합을 구하는 프로그램을 작성하시오. 즉, N개의 수가 주어졌을 때, 길이가 K인 연속 부분 수열 N-K+1개의 중앙값의 합을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N ≤ 250,000, 1 ≤ K ≤ 5,000, K ≤ N)

둘째 줄부터 N개 줄에 측정한 온도가 순서대로 주어진다. 온도는 0보다 크거나 같고, 65535보다 작거나 같은 정수이다.

출력

길이가 K인 모든 연속 부분 수열의 중앙값의 합을 출력한다.

풀이

#include <iostream>
#include <algorithm>

class smt {
public:
	smt() { }

	int update(int node, int start, int end, int index, int num) {
		if (index < start || index > end) {
			return mTree[node];
		}
		else if (start == end) {
			mTree[node] += num;
			return mTree[node];
		}

		int mid = (start + end) / 2;

		mTree[node] = update(node * 2, start, mid, index, num) + update(node * 2 + 1, mid + 1, end, index, num);

		return mTree[node];
	}

	int query(int node, int start, int end, int num) {
		if (start == end) {
			return start;
		}

		int mid = (start + end) / 2;

		if (mTree[node * 2] >= num) {
			return query(node * 2, start, mid, num);
		}
		else {
			return query(node * 2 + 1, mid + 1, end, num - mTree[node * 2]);
		}
	}

private:
	int mTree[65536 * 4];
};

int arr[250001] = { 0, };
smt seg;

int main() {
	std::ios::sync_with_stdio(false);
	std::cout.tie(NULL);
	std::cin.tie(NULL);

	int N;
	int K;
	std::cin >> N >> K;

	for (int i = 0; i < N; i++) {
		std::cin >> arr[i];
	}

	for (int i = 0; i < K; i++) {
		seg.update(1, 0, 65536, arr[i], 1);
	}

	long long ans = seg.query(1, 0, 65536, (K + 1) / 2);

	for (int i = K; i < N; i++) {
		seg.update(1, 0, 65536, arr[i], 1);
		seg.update(1, 0, 65536, arr[i - K], -1);

		int mid = seg.query(1, 0, 65536, (K + 1) / 2);

		ans += mid;
	}

	std::cout << ans << "\n";

	return 0;
}