ICPC Regionals에 나온 문제입니다.

IT 회사의 직원들이 \(1\)번부터 \(N\)번까지 있습니다. 각 직원들은 고유의 역량 \(P_i\) 갖습니다.

임의의 과제에 대해, \(L\)번부터 \(R\)번까지, 역량이 \(A\) 이상 \(B\) 이하인 직원들을 배정할 수 있습니다. 단, 배정되는 직원들의 역량의 평균은 \(S\) 이상이 되어야 합니다. \(Q\)개의 과제가 주어집니다. 각 과제에 대해 최대한 많은 직원들을 배정하고자 할 때, 인원의 최대값을 구하면 됩니다.

평균이 \(S\) 이상이 되게 하면서 최대한 많은 사람들을 넣고자 합니다. 그리디하게, 역량이 높은 직원부터 배정하면서 역량이 \(S\) 미만이 되는 시점에 배정을 중단하고 인원수를 찾으면 됩니다.

하지만 이 방식은 과제가 1개일 때 유효한 방식입니다. 여러 개의 쿼리를 처리하기 위해선 최적화가 필요합니다.

이 문제에 접근하는 방법은 3가지 이상이 있습니다.

  • Mo’s
  • 퍼시스턴트 세그먼트 트리 + 이분 탐색
  • 펜윅 트리 + 병렬 이분 탐색

Mo’s Algorithm

이 방식은 구간 쿼리를 재배열하여 중복되는 구간 연산 값들을 재활용합니다. \(N\) 크기와 시간 제한을 생각했을 때 좀 빡빡하긴 합니다. 이 방식을 적용한 정해 코드를 찾았으나 직접 작성한 것은 아니므로 여기서 다루지는 않을 것입니다.

이분 탐색

어떤 \(L\)번부터 \(R\)번까지의 직원들만 고려할 때, 역량의 평균이 \(S\) 이상이 되는 최대값을 찾아봅시다.

역량의 상한 \(B\)는 고정합니다. 자명하게 역량이 큰 순서대로 넣는 것이 유리하기 때문입니다. 가상의 하한 \(m\)을 잡아서, 역량이 \(m\)부터 \(B\)까지인 직원들의 평균을 내봅니다.

평균이 \(S\) 이상이라면, 역량이 \(m\) 미만인 직원을 몇 명 더 넣을 수 있습니다. \(S\) 미만이라면, 역량이 작은 몇 명을 빼야한다는 뜻입니다.

결정 문제로 치환할 수 있습니다. 임의의 구간의 역량 평균을 빠르게 구할 수 있는 방식을 쓰면 되는데, 세그먼트 트리 등을 활용해 \(O(\log^2 P)\)로 답을 구할 수 있습니다.

Persistent Segment Tree

이 문제를 처음 해결했을 때 고른 방식입니다.

위에서 제시한 이분 탐색 방식에서 병목이 되는 부분은 \(L\)번부터 \(R\)번까지만 고려한다는 점입니다. \(Q\)번의 쿼리를 처리해야 하므로 매 쿼리마다 구간 길이 \(R-L+1\)를 세그먼트 트리에 업데이트하는 방식으로는 시간 내에 해결할 수 없습니다.

PST는 이런 경우 적용할 수 있는 자료구조입니다. 매 직원의 역량을 세그먼트 트리에 넣으면서 그 기록을 남길 수 있습니다. \(L\)번부터 \(R\)번까지의 업데이트 기록은 \(R\)번 기록에서 \(L-1\)번을 빼서 구할 수 있습니다.

struct Result {
	int cnt;
	ll sum;
	Result operator+(const Result& o) const { return { cnt + o.cnt, sum + o.sum }; }
};

struct Node { int l, r, cnt; ll sum; } t[LEN * 22]; // 공간복잡도가 N log N이므로 충분히 할당합니다.
int cur = 0, ptr[LEN];

struct PersistentSegTree {
	void update(int k, int x, int d) { update(ptr[k - 1], ptr[k] = ++cur, 1, LEN, x, d); }
	void update(int p, int i, int s, int e, int x, int d) {
		if (s == e) { t[i].cnt = t[p].cnt + 1; t[i].sum = t[p].sum + d; return; }
		int m = s + e >> 1;
		if (x <= m) {
			t[i].l = ++cur; t[i].r = t[p].r;
			update(t[p].l, t[i].l, s, m, x, d);
		}
		else {
			t[i].r = ++cur; t[i].l = t[p].l;
			update(t[p].r, t[i].r, m + 1, e, x, d);
		}
		t[i].cnt = t[t[i].l].cnt + t[t[i].r].cnt;
		t[i].sum = t[t[i].l].sum + t[t[i].r].sum;
	}
	// p + 1 ~ i 시점의 구간 기록 쿼리
    Result get(int p, int i, int l, int r) { return get(ptr[p], ptr[i], l, r, 1, LEN); }
    Result get(int p, int i, int l, int r, int s, int e) {
        if (r < s || e < l) return {0 , 0};
        if (l <= s && e <= r) return { t[i].cnt - t[p].cnt, t[i].sum - t[p].sum };
        int m = s + e >> 1;
        return get(t[p].l, t[i].l, l, r, s, m) + get(t[p].r, t[i].r, l, r, m + 1, e);
    }
} pst;

// ...

ll l = A, r = B, m, cnt, ret = B + 1, sum, ans = 0;

while (l <= r) {
	m = l + r >> 1;
	Result cur = pst.get(L - 1, R, m, B); // L ~ R번까지, 역량이 m ~ B인 직원들을 구합니다.
	sum = cur.sum;
	cnt = cur.cnt;
	if (sum >= cnt * S) {
		ret = std::min(ret, m);
		r = m - 1;
	}
	else l = m + 1;
}

세그먼트 트리 업데이트와 쿼리 처리까지 다 해서 \(O(N \log P + Q \log^2 P)\)으로 처리할 수 있습니다. 여기까지만 해도 사실 문제를 해결하기는 충분합니다.

병렬 이분 탐색

하지만 PST를 굳이 쓸 필요가 있을지는 다시 따져볼 필요가 있습니다.

\(Q\)개의 쿼리가 각각의 \(m\) 값을 가지고 있다고 합시다. 이제 \(1\)번 직원부터 펜윅트리에 넣겠습니다. \(L-1\)번 직원을 넣은 시점에 \(m\)~\(B\) 구간 쿼리로 값 \(x_l\)을 미리 구해둡니다. \(R\)번 직원을 넣은 시점에 다시 한 번 구간 쿼리로 값 \(x_r\)을 구합니다. \(L\)번부터 \(R\)번 직원들에 대한 구간 쿼리는 \(x_r - x_l\) 값을 구하면 됩니다. 모든 \(Q\)개의 쿼리에 대해, 구간 업데이트를 스위핑하듯 처리할 수 있습니다. 이렇게 해서 뭘 얻을 수 있을까요?

스위핑이 끝난 후, 각각의 쿼리애 대해 \(m\) 값을 업데이트할 겁니다. \(m\)값을 찾는 과정은 결국 이분탐색이므로, 업데이트는 많아야 \(\log N\)번 이루어집니다. 스위핑을 \(\log P\)번 하고나면, 모든 쿼리는 \(m\)값을 찾게 됩니다.

발상은 상당히 어렵지만, 구현은 꽤 간단합니다.

#include <iostream>
#include <algorithm>
#include <cstring>
#include <vector>

typedef long long ll;
const int LEN = 200'010;

struct Fenwick {
	ll t[LEN];
	ll sum(int i) {
		ll ret = 0;
		while (i > 0) {
			ret += t[i];
			i -= i & -i;
		}
		return ret;
	}
	void update(int i, ll d) {
		while (i < LEN) {
			t[i] += d;
			i += i & -i;
		}
	}
} sum, cnt;

struct Query { 
	int l, r;
	ll sum, cnt;
} q[LEN];

struct Project {
	int l, r;
	int a, b;
	ll s;
} p[LEN];

int N, Q, P[LEN];

struct Event {
	int t;
	int i, d;
	bool operator<(const Event& r) const { return t < r.t; }
} e[LEN * 2];

int main() {
	std::cin.tie(0)->sync_with_stdio(0);
	std::cin >> N;
	for (int i = 1; i <= N; ++i) std::cin >> P[i];

	std::vector<int> cur;

	std::cin >> Q;
	for (int i = 0; i < Q; ++i) {
		std::cin >> p[i].l >> p[i].r >> p[i].a >> p[i].b >> p[i].s;
		e[i << 1] = { p[i].l - 1, i, -1 };
		e[i << 1 | 1] = { p[i].r, i, +1 };
		q[i].l = p[i].a;
		q[i].r = p[i].b;
		cur.push_back(i);
	}
	std::sort(e, e + Q * 2);

	while (cur.size()) { // 처리할 쿼리가 남아있을 동안
		memset(sum.t, 0, sizeof sum.t);
		memset(cnt.t, 0, sizeof cnt.t);
		int j = 0;
		while (j < Q * 2 && e[j].t == 0) {
			q[e[j].i].sum = q[e[j].i].cnt = 0;
			++j;
		}

		for (int i = 1; i <= N; ++i) { // 스위핑
			sum.update(P[i], P[i]);
			cnt.update(P[i], 1);

			while (j < Q * 2 && e[j].t == i) {
				int x = e[j].i;
				int m = q[x].l + q[x].r >> 1;
				if (e[j].d == 1) { // R 시점
					q[x].cnt += cnt.sum(p[x].b) - cnt.sum(m - 1);
					q[x].sum += sum.sum(p[x].b) - sum.sum(m - 1);
				}
				else { // L - 1 시점
					q[x].cnt = cnt.sum(m - 1) - cnt.sum(p[x].b);
					q[x].sum = sum.sum(m - 1) - sum.sum(p[x].b);
				}
				++j;
			}
		}

		std::vector<int> nxt;
		for (const int& i : cur) {
			int m = q[i].l + q[i].r >> 1;			
			if (q[i].sum >= q[i].cnt * p[i].s) q[i].r = m;
			else q[i].l = m + 1;

			if (q[i].l < q[i].r) nxt.push_back(i);
		}
		cur.swap(nxt);
	}

	for (int i = 0; i < Q; ++i) {
		int ans = q[i].sum >= q[i].cnt * p[i].s ? q[i].cnt : 0;
		if (ans && q[i].r > p[i].a) {
			int d = (q[i].sum - p[i].s * q[i].cnt) / (p[i].s - q[i].r + 1);
			ans += d;
		}
		std::cout << ans << '\n';
	}
}

병렬 이분 탐색 방식은 PST 방식에 비해 몇 가지 장점이 있습니다.

  1. 공간복잡도가 작다.
  • PST가 비교적 적은 메모리로 모든 업데이트 과정을 보존한다고는 해도, 여전히 큰 공간을 필요로 합니다.
  1. 비교적 빠르다.
  • 큰 공간을 차지한다는 것을 접근에 필요한 시간 또한 길어짐을 의미합니다. PST는 다소 느립니다.
  • 재귀 세그먼트 트리 대신 시간복잡도 면에서 유리한 Fenwick 트리를 쓸 수도 있습니다.
  1. 구현이 간단하다.

대신, 이 방법은 각 쿼리가 독립이라는 조건이 붙습니다. 온라인 쿼리의 경우에는 적용할 수 없습니다.

병렬이분탐색까지는 아니지만, 이 글을 읽은 후 최적화를 고려할 수 있는 문제는 다음과 같습니다.