[백준] 9426(세그먼트 트리)

백준 9426문제는 N개의 원소중에서 K개의 부분집합원소의 중앙값들의 합을 구하는 문제입니다.

이 문제는 세그먼트 트리를 이용해서 풀었습니다.

세그먼트 트리의 각노드에는 노드들의 갯수를 저장하고있습니다. (리프 노드에는 해당 노드가 세그먼트 트리에있는지 없는지를 표시)

업데이트시 1, -1로 원소를 추가 삭제 가능합니다.

쿼리함수를 이용해서 서브트리에 대한 중간값에 해당하는 값을 찾을수있습니다.

세그먼트 트리는 보통 포화 이진트리이므로 값이 N개일 떄 높이는 O(logN)입니다.

즉 하나의 중간값을 찾는데 걸리는시간은 O(logN)이 걸립니다.

그러므로 세그먼트 트리를 이용한 최악의 경우 시간 복잡도는 약 O(NlogN)입니다.

Worst Case time Complexity = O(NlogN)

include <iostream>
#include <algorithm>
// 세그먼트 트리사용
using namespace std;
typedef long long ll;

#define MAX_N 250001
#define MAX 65536

ll seg[4 * MAX_N], arr[MAX_N], N, K, sum;// 위치를 기준으로 값을 저장

ll update(ll pos, ll val, ll node, ll x, ll y) { // 노드와 노드의 범위
	if (pos < x || pos > y) return seg[node];
	if (x == y) return seg[node] += val;
	ll mid = (x + y) >> 1; // 한비트를 당김 즉 (x + y) / 2와 같음
	return seg[node] = update(pos, val, node * 2, x, mid) + update(pos, val, node * 2 + 1, mid + 1, y);
}

ll query(ll pos, ll node, ll x, ll y) {
	if (x == y) return x; // 해당값을 찾았을경우 그값을 리턴

	ll mid = (x + y) >> 1;
	if (pos <= seg[2 * node]) return query(pos, 2 * node, x, mid); // 상대적인 위치를 찾음
	else return query(pos - seg[2 * node], 2 * node + 1, mid + 1, y); 
}

int main() {
	freopen("test.txt", "r", stdin);
	std::ios::sync_with_stdio(false);
	cin >> N >> K;
	ll answer = 0;

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

	for (int i = 0; i < K; ++i) {
		update(arr[i], 1, 1, 0, MAX_N - 1);
	}
	
	ll mid = (K + 1) / 2;
	for (int i = 0; i < N - K + 1; ++i) {
		answer += query(mid, 1, 0, MAX_N - 1);
		update(arr[i], -1, 1, 0, MAX_N - 1);
		if (i == N - K) break; // 업데이트 할필요가 없음
		update(arr[i + K], 1, 1, 0, MAX_N - 1);
	}

	cout << answer << endl;
	return 0;
}