일반화 segment tree 구현 - 3. 업데이트와 쿼리 구현

저번 글에 이어서 이번에는 세그먼트 트리의 업데이트와 쿼리 부분을 구현하여 세그먼트 트리를 완성해보았다.

1. 업데이트

업데이트의 구현의 기본적으로 트리 초기화와 크게 다르지 않다. 다만 특정 값을 더하거나 곱한다던지 아니면 아예 다른 값으로 대체한다던지 업데이트의 방식이 그때그때 다르므로 업데이트하는 방식 또한 함수 포인터로 받도록 했다.

void update(int idx,vector<T> val,T (*func)(T, vector<T>)) {
		upd(1,s,e,idx,val,func);
} //사용자가 실제로 사용할 함수

T upd(int node,int start,int end,int idx,vector<T> &val,T (*func)(T, vector<T>)) {
		if(idx<start||idx>end) return tree[node];
		if(start==end) return tree[node]=func(tree[node],val);
		int mid=(start+end)/2;
		return tree[node]=merge(upd(node*2,start,mid,idx,val,func),upd(node*2+1,mid+1,end,idx,val,func));
} //구조체 내부에서 업데이트를 위해 사용할 함수, idx는 업데이트할 인덱스 val은 업데이트에 사용할 값들

2. 쿼리

쿼리도 기존의 세그먼트 트리의 쿼리 해결 방식과 크게 다르지는 않지만 범위에 포함되지 않을시 리턴할 값을 정하는 것이 약간 문제가 됐다.

처음에는 NULL을 사용하여 왼쪽 노드와 오른쪽 노드 중 하나의 리턴값이 NULL일 경우 반대쪽 노드의 값만 리턴하도록 코딩을 해보았는데 제대로 작동은 됐으나 경고 메시지가 나왔다.

다른 언어에 있는 undefine이나 none같은 것이 c++에 존재하는지 찾아보았지만 c++에서는 존재하지 않아 직접 리턴할 값을 입력 받도록 했다.

T query(int left,int right,T null) {
        return que(left,right,1,s,e,null);
} //사용자가 실제로 사용할 함수

T que(int left,int right,int node,int start,int end,T null) {
        if(right<start||left>end) return null;
        if(start>=left&&end<=right) return tree[node];
        int mid=(start+end)/2;
        return merge(que(left,right,node*2,start,mid,null),que(left,right,node*2+1,mid+1,end,null));
} //구조체 내부에서 쿼리를 위해 사용할 함수, null은 범위에 포함되지 않을 시 리턴할 값

이로써 일반화 프로그래밍으로 세그먼트 트리를 완성했다. 아래는 완성된 전체 코드이다.

template<typename T> struct segment_tree {
	
	T (*merge)(T, T);
	int s,e,N;
	vector<T> tree;
	
	segment_tree(vector<T> vec,int start,int end,T (*func)(T, T)) {
		s=start;
		e=end;
		N=e-s+1;
		merge=func;
		tree.resize(4*N+1);
		init(vec,1,start,end);
	}
	
	T query(int left,int right,T null) {
		return que(left,right,1,s,e,null);
	}
	
	void update(int idx,vector<T> val,T (*func)(T, vector<T>)) {
		upd(1,s,e,idx,val,func);
	}
	
	private: //사용자가 직접 접근하지 않는 함수는 private처리
	
	T init(vector<T> &vec,int node,int start,int end) {
		if(start==end) return tree[node]=vec[start];
		int mid=(start+end)/2;
		return tree[node]=merge(init(vec,node*2,start,mid),init(vec,node*2+1,mid+1,end));
	}
	
	T upd(int node,int start,int end,int idx,vector<T> &val,T (*func)(T, vector<T>)) {
		if(idx<start||idx>end) return tree[node];
		if(start==end) return tree[node]=func(tree[node],val);
		int mid=(start+end)/2;
		return tree[node]=merge(upd(node*2,start,mid,idx,val,func),upd(node*2+1,mid+1,end,idx,val,func));
	}
	
	T que(int left,int right,int node,int start,int end,T null) {
		if(right<start||left>end) return null;
		if(start>=left&&end<=right) return tree[node];
		int mid=(start+end)/2;
		return merge(que(left,right,node*2,start,mid,null),que(left,right,node*2+1,mid+1,end,null));
	}
};

다음 글에서는 마지막으로 직접 만든 세그먼트 트리로 몇가지 알고리즘 문제를 풀어보겠다.