ARC033C データ構造

シェアする

  • このエントリーをはてなブックマークに追加

こんにちは。

昨日は一昨日解いた問題の解説を書くだけになってしまって問題を解けませんでしたが、本日は解きました。

というか、昨日のうちに解こうと思ってた問題でミスに悩まされ、今朝やっと解けたという感じです。

さて、今回解いたのはARC033のC問題「データ構造」です。

数の集合(最初は空集合)に「Xを追加せよ」と「小さい方からX番目の数を取り出せ」という命令を大量にこなしていくというものです。

パッと見で「二分探索木とかかな?」と思ったのですが、下からX番目を見つける機能が無いのでダメでした。

解説を読むとどうやら「セグメント木」を使うようです。RMQ(Range Minimum Query)で使うやつです。

セグメント木の使い方さえ分かればそれほど難しい問題ではありません。

ということで考えていきましょう。

セグメント木の機能

今回の問題は「数列にXを追加する機能」と「下からX番目を返す機能」があるデータ構造が必要になりますが。後者の機能は「ある要素より小さい要素が何個あるか」が分かれば、二分探索を用いて高速に実現することができます。

この機能は、要素数200,001の数列(int dat[200001]とする)を用意し、「Xを追加」が来ればdat[X]を1増やし、「要素dat[i]より小さい要素が何個あるか」を計算するときは「dat[0]+dat[1]+...+dat[i-1]」を行うことで(ソート無しで)実現できます。

ただし、追加操作はO(1)で行えますが、追加するXの最大値をN(最大で200,001)とすると、区間の和を求める計算にO(N)の時間が掛かり、クエリがQ個(最大で200,001)あるので、単純な実装ではO(NQ)の時間が掛かってしまいます。

N,Q=200,001が最大なので、O(NQ)では制限時間2秒を超えてしまいます。

そこで、区間に対する特定の処理(この問題では区間の和)をO(logN)で行えるデータ構造「セグメント木」を利用します。(データ追加・変更もO(logN)時間かかります)

セグメント木の詳細については蟻本(p.153)や他のサイトを確認していただくとして、ここではざっくりとだけ紹介します。

※「蟻本」とは「プログラミングコンテストチャレンジブック」という競技プログラミング向けテクニック集のような書籍のことです。表紙に蟻の絵があるので蟻本。競プロ初心者~中級者で強くなりたい人は必読です。

セグメント木とは、「数列の区間に対する特定の処理」を扱うためのデータ構造です。

与えられた数列を2分木の葉ノードとして保持し、各ノードの親ノードは「2つの子ノードの値の最小値や和など」を持ちます。RMQの問題なら最小値、今回のような区間の和を扱う問題なら和を持ちます。

セグメント木が有効なのは、区間に対する処理が「区間[a,b)に対する処理の結果が、[a,c)ととして、dat[0]の子がdat[1]とdat[2]、dat[1]の子がdat[3]とdat[4]、dat[2]の子がdat[5]とdat[6]、……

というように管理すると、

・dat[i]の子はdat[2*i+1]とdat[2*i+1]

・dat[i]の親はdat[(i-1)/2]

・葉ノードの数をnとするとdat[]の要素数=ノード数は2*n-1

となります。

では、これを実装してみます。

セグメント木の実装

今朝提出してACしたプログラムのセグメント木部分は以下のようになっています。

struct SumSegTree {
	int n, height;
	vector<int> dat;

	// 初期化(_nは最大要素数)
	SumSegTree(int _n) {
		n = 1;
		height = 1;
		while (n < _n) {
			n *= 2;
			height++;
		}
		dat = vector<int>(2 * n - 1);
	}

	// 場所i(0-indexed)にxを足す
	void add(int i, int x) {
		i += n - 1; // i番目の葉ノードへ
		dat[i] += x;
		while (i > 0) { // 下から上がっていく
			i = (i - 1) / 2;
			dat[i] += x;
		}
	}

	// 区間[a,b)の総和。ノードk=[l,r)に着目している。
	int _sum(int a, int b, int k, int l, int r) {
		if (r <= a || b <= l)return 0;	// 交差しない
		if (a <= l && r <= b)return dat[k];	// a,l,r,bの順で完全に含まれる
		else { 
			int s1 = _sum(a, b, 2 * k + 1, l, (l + r) / 2);	// 左の子
			int s2 = _sum(a, b, 2 * k + 2, (l + r) / 2, r);	// 右の子
			return s1 + s2;
		}
	}

	// 内部クエリ_sum()を呼び出す
	int sum(int a, int b) {
		return _sum(a, b, 0, 0, n);
	}
};

要素数N以上のセグメント木は SumSegTree(N) で宣言します。すると、Nを超える最小の"2の累乗"である数が変数nに格納され、葉ノード数nのセグメント木が用意されます。

区間和を計算する関数 _sum() は再帰関数になっています。条件分岐が複雑なので注意してください。

これを用いて、問題を解くプログラムを作成します。

解答作成

今回行うべき処理は、「T=1のとき、add(X,1)を行う」「T=2のとき、sum(1,t)=Xになるtを見つけて出力し、add(t,-1)を行う」です。

後者の探索は二分探索を用いて行います。

提出したC++コードがこちら。

#include <iostream>
#include <cstdio>
#include <string>
#include <algorithm>
#include <vector>

#define repl(i,a,b) for(int i=(int)(a);i<(int)(b);i++)
#define rep(i,n) repl(i,0,n)

typedef long long ll;

using namespace std;

const int MOD = 1000000007;

struct SumSegTree {
	/* 省略。上記のコードです */
}; 

int Q, N = 200001;

int main() { cin >> Q;
	SumSegTree segt(N);
	rep(i, Q) {
		int T, X;
		cin >> T >> X;
		if (T == 1) {
			segt.add(X, 1);
		} else {
			int ans = 0, r = N;
			// 二分探索
			while (r - ans > 1) {
				int mid = (ans + r) / 2;
				if (segt.sum(1, mid) < X) {
					ans = mid;
				}
				else {
					r = mid;
				}
			}
			cout << ans << endl;
			segt.add(ans, -1);
		}
	}
	return 0;
}

今回はセグメント木で実装しましたが、この問題で必要な処理は区間の和なので、BIT(Binary Indexed Tree)を用いることができます。

前回解いた問題でも使いましたが、「dat[a]からdat[b]までの和」は「dat[0]からdat[i]までの和」を高速に計算できれば十分です。

始めから途中までの区間和しか計算しないことを前提にすれば、セグメント木の中で不要なノードが出てきます。

詳しくは蟻本のp.159などを参照してください。

それでは、本日はここまで。

BIT(Binary Indexed Tree)での実装

解説は蟻本等に任せますが、BITでの実装がすごく簡単だったので追記しておきます。

#include <iostream>
#include <cstdio>
#include <string>
#include <algorithm>
#include <vector>

#define repl(i,a,b) for(int i=(int)(a);i<(int)(b);i++)
#define rep(i,n) repl(i,0,n)

typedef long long ll;

using namespace std;

const int MOD = 1000000007;

struct BIT{
	int n, height;
	vector<int> dat;

	BIT(int _n) {
		n = 1;
		height = 1;
		while (n < _n) {
			n *= 2;
			height++;
		}
		dat = vector<int>(n + 1);
	}

	// [0,i)までの和
	int sum(int i) {
		int s = 0;
		while (i > 0) {
			s += dat[i];
			i -= i & -i;	// i & -i は i の最後の1ビット
		}
		return s;
	}

	void add(int i, int x) {
		i++;	// 0-indexedに変更
		while (i <= n) {
			dat[i] += x;
			i += i & -i;
		}
	}
};

int Q, N = 200001;

int main() {
	cin >> Q;
	BIT bit(N);
	rep(i, Q) {
		int T, X;
		cin >> T >> X;
		if (T == 1) {
			bit.add(X, 1);
		} else {
			int ans = 0, r = N;
			// 二分探索
			while (r - ans > 1) {
				int mid = (ans + r) / 2;
				if (bit.sum(mid) < X) {
					ans = mid;
				}
				else {
					r = mid;
				}
			}
			cout << ans << endl;
			bit.add(ans, -1);
		}
	}
	return 0;
}
スポンサーリンク
レクタングル(大)
レクタングル(大)

シェアする

  • このエントリーをはてなブックマークに追加

フォローする