セグメント木を使ってみた

データ構造 Codeforces C++

Competitive Programming Advent Calendar Div201218日目の記事、kagamizさんのIntroduction to Segment Treeを見て、セグメント木を実装してみました。

解いたのはCodeforcesのC. Circular RMQ。久々にC++で実装しました。スクリプト言語でやろうかとも思ったのですが、最近スクリプト言語で大きいデータを扱うのが怖いです・・・。

#include <cstdio>
#include <algorithm>
#include <string.h>
#include <limits.h>
#include <iostream>
using namespace std;

#define MAX_SIZE 200000 * 4
typedef long long ll;

//Starry Sky Tree
//http://kagamiz.hatenablog.com/entry/2012/12/18/220849

ll segMin[MAX_SIZE], segAdd[MAX_SIZE];

//区間[a, b)に値xを加算する.
void add(int a, int b, int x, int k, int l, int r)
{
	if (r <= a || b <= l) return; //もし交差しない区間であれば終える.
	
	if (a <= l && r <= b){ //もし今みている区間[l, r)が[a, b)に完全に内包されていれば
		segAdd[k] += x;  //区間[l, r)にkを加算する.
		while (k){
			k = (k - 1) / 2;
                        //親の区間の最小値は, 子の区間の最小値 + 自分に一様に加算されている値 である.一様に加算される値は更新しなくて良い.
			segMin[k] = min(segMin[k * 2 + 1] + segAdd[k * 2 + 1], segMin[k * 2 + 2] + segAdd[k * 2 + 2]);
		}
		return;
	}
	
	add(a, b, x, k * 2 + 1, l, (l + r) / 2); //子の区間に(必要があれば)xを加算する.
	add(a, b, x, k * 2 + 2, (l + r) / 2, r); //〃
}

ll getMin(int a, int b, int k, int l, int r)
{
	if (r <= a || b <= l) return (LLONG_MAX);
	
	if (a <= l && r <= b) return (segMin[k] + segAdd[k]); //完全に内包されていれば,その区間の最小値を返す.
	
	ll left = getMin(a, b, k * 2 + 1, l, (l + r) / 2); //子の区間の最小値を求める.
	ll right = getMin(a, b, k * 2 + 2, (l + r) / 2, r); //子の区間の最小値を求める
	
	return (min(left, right) + segAdd[k]); //親の区間の最小値は, 子の区間の最小値 + 自分に一様に加算されている値 である (大切なので2回書きました!!)
	
}

int main(){
	char ss[1024];
	int n, m;
	scanf("%d", &n);
	for(int i = 0; i < n; i++){
	  int t;
	  scanf("%d", &t);
	  add(i, i + 1, t, 0, 0, n);
	} 
	scanf("%d", &m);
	fgets(ss, sizeof(ss), stdin);
	for(int i = 0; i < m; i++){
	  int args[3];
	  fgets(ss, sizeof(ss), stdin);
	  char *tok;
	  int count = 0;
	  tok = strtok( ss, " " );
	  while( tok != NULL ){
	    args[count++] = atoi(tok);
	    tok = strtok( NULL, " " );  /* 2回目以降 */
	  }
	  if(count == 2){
	    int lf = args[0];
	    int rg = args[1];
	    ll ans = 0;
	    if(rg < lf) ans = min(getMin(lf, n, 0, 0, n), getMin(0, rg + 1, 0, 0, n));
	    else ans = getMin(lf, rg + 1, 0, 0, n);
	    cout << ans << "\n";
	  }
	  else{
	    int lf = args[0];
	    int rg = args[1];
	    int v = args[2];
	    if(rg < lf){
	      add(lf, n, v, 0, 0, n);
	      add(0, rg + 1, v, 0, 0, n);
	    }
	    else add(lf, rg + 1, v, 0, 0, n);
	  }
	} 
	
	return 0;
}


セグメント木の部分は丸コピーです。
この問題で難しかったのはStarry Sky木（セグメント木の改良版）の実装とsplit関数の実現です。
C++には意外にもsplit関数がなく、代わりにstrtok関数を使うことで似たような機能を実装できます（参考 : http://simd.jugem.jp/?eid=136、http://kagamiz.hatenablog.com/entry/2012/12/11/233004）。

MAX_SIZEはn * 2あれば十分のはずなのですが、なぜかRuntime Errorになってしまったので、n * 4にしてみたら通りました。多分addとgetMinでちょっと余分に見に行ってしまっているのでしょうか。