취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 13519번 문제인 트리와 쿼리 10이다.
문제는 아래 링크를 확인하자.

https://www.acmicpc.net/problem/13519 

트리 위에서의 경로 내 최대 연속합 쿼리와 경로 업데이트 쿼리를 빠르게 처리하는 문제이다.

배열 위에서의 범위 내 최대 연속합을 구하는 방법은 잘 알려져있다는 점을 상기하면, 해당 문제를 배열 위에서 해결하는 세그먼트트리와 HLD(heavy-light decomposition)를 이용해 문제를 해결할 수 있다는 것을 알아낼 수 있다.

구현량이 매우 많으므로 실수 없이 잘 구현하자.

아래는 제출한 소스코드이다.

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef long long ll;

int N, Q;

struct node {
	ll val, lval, rval, tval;
	node() {};
	node(ll val, ll lval, ll rval, ll tval) {
		this->val = val, this->lval = lval, this->rval = rval, this->tval = tval;
	}
};

node nodemerge(node& nodeL, node& nodeR) {
	ll retV = max(nodeL.rval + nodeR.lval, max(nodeL.val, nodeR.val));
	ll retL = max(nodeL.lval, nodeL.tval + nodeR.lval);
	ll retR = max(nodeR.rval, nodeR.tval + nodeL.rval);
	ll retT = nodeL.tval + nodeR.tval;
	return node(retV, retL, retR, retT);
}

ll prearr[100001];
ll arr[100001];
node seg[262145];
ll lazy[262145];
node init(int L, int R, int sI) {
	if (L == R) {
		ll& tmp = arr[L];
		if (tmp < 0) return seg[sI] = node(0, 0, 0, tmp);
		else return seg[sI] = node(tmp, tmp, tmp, tmp);
	}

	node nodeL = init(L, (L + R) / 2, sI * 2), nodeR = init((L + R) / 2 + 1, R, sI * 2 + 1);
	return seg[sI] = nodemerge(nodeL, nodeR);
}
void propagate(int L, int R, int sI) {
	ll& tmp = lazy[sI]; tmp -= 65536;
	if (tmp < 0) seg[sI] = node(0, 0, 0, tmp * (R - L + 1));
	else seg[sI] = node(tmp * (R - L + 1), tmp * (R - L + 1), tmp * (R - L + 1), tmp * (R - L + 1));
	if (L < R) lazy[sI * 2] = lazy[sI * 2 + 1] = tmp + 65536;
	lazy[sI] = 0;
}
node upd(int L, int R, int qL, int qR, ll qVal, int sI) {
	if (lazy[sI]) propagate(L, R, sI);
	if (R < qL || qR < L) return seg[sI];
	if (qL <= L && R <= qR) {
		lazy[sI] = qVal + 65536;
		propagate(L, R, sI);
		return seg[sI];
	}
	node nodeL = upd(L, (L + R) / 2, qL, qR, qVal, sI * 2), nodeR = upd((L + R) / 2 + 1, R, qL, qR, qVal, sI * 2 + 1);
	return seg[sI] = nodemerge(nodeL, nodeR);
}
node qry(int L, int R, int qL, int qR, int sI) {
	if (lazy[sI]) propagate(L, R, sI);
	if (R < qL || qR < L) return node(0, 0, 0, 0);
	if (qL <= L && R <= qR) return seg[sI];

	node nodeL = qry(L, (L + R) / 2, qL, qR, sI * 2), nodeR = qry((L + R) / 2 + 1, R, qL, qR, sI * 2 + 1);
	return nodemerge(nodeL, nodeR);
}

vector<int> G[100001];
int childcnt[100001];
int heavychild[100001];

int dfs1(int cur, int par) {
	int ret = 0;
	int heavycnt = 0;
	for (auto node : G[cur]) {
		if (node == par) continue;
		int tmp = dfs1(node, cur);
		if (tmp > heavycnt) heavychild[cur] = node, heavycnt = tmp;
		ret += tmp;
	}
	childcnt[cur] = ret;
	return ret + 1;
}

int nodeidx = 1, chainidx = 1;
int dfi[100001];
//아래 배열들은 dfi기준
int chainnum[100001]; // 현재 노드는 몇번 체인에 속해있는가
int chainnth[100001]; // 현재 노드는 속해있는 체인의 몇번째(0-based) 원소인가
int chainlevel[100001]; // 현재 노드가 포함된 체인의 깊이는 몇인가
int chainparent[100001]; // 현재 노드가 포함된 체인은 어디(dfi)에 붙어있는가

void dfs2(int cur, int par, int cpar, int nth, int lv) {
	int idx = dfi[cur] = nodeidx++;
	arr[idx] = prearr[cur];
	chainnum[idx] = chainidx;
	chainnth[idx] = nth;
	chainlevel[idx] = lv;
	chainparent[idx] = cpar;

	int& hchild = heavychild[cur];
	if (hchild) dfs2(hchild, cur, cpar, nth + 1, lv);

	for (auto node : G[cur]) {
		if (node == par || node == hchild) continue;
		chainidx++;
		dfs2(node, cur, idx, 0, lv + 1);
	}
}

void solve() {
	int q; cin >> q;
	if (q == 1) {
		int u, v; cin >> u >> v;
		u = dfi[u], v = dfi[v];
		if (chainlevel[u] > chainlevel[v]) swap(u, v);
		node nodeU = node(0, 0, 0, 0), nodeV = node(0, 0, 0, 0);
		while (chainlevel[u] < chainlevel[v]) {
			node nodeQ = qry(1, N, v - chainnth[v], v, 1);
			nodeV = nodemerge(nodeQ, nodeV);
			v = chainparent[v];
		}
		while (chainnum[u] != chainnum[v]) {
			node nodeQ = qry(1, N, u - chainnth[u], u, 1);
			nodeU = nodemerge(nodeQ, nodeU);
			nodeQ = qry(1, N, v - chainnth[v], v, 1);
			nodeV = nodemerge(nodeQ, nodeV);
			u = chainparent[u];
			v = chainparent[v];
		}
		if (u > v) swap(u, v), swap(nodeU, nodeV);
		node nodeQ = qry(1, N, u, v, 1);
		swap(nodeU.lval, nodeU.rval);
		nodeU = nodemerge(nodeU, nodeQ);
		nodeU = nodemerge(nodeU, nodeV);
		cout << nodeU.val << '\n';
	}
	else {
		int u, v; ll qVal; cin >> u >> v >> qVal;
		u = dfi[u], v = dfi[v];
		if (chainlevel[u] > chainlevel[v]) swap(u, v);
		while (chainlevel[u] < chainlevel[v]) {
			upd(1, N, v - chainnth[v], v, qVal, 1);
			v = chainparent[v];
		}
		while (chainnum[u] != chainnum[v]) {
			upd(1, N, u - chainnth[u], u, qVal, 1);
			upd(1, N, v - chainnth[v], v, qVal, 1);
			u = chainparent[u];
			v = chainparent[v];
		}
		if (u > v) swap(u, v);
		upd(1, N, u, v, qVal, 1);
	}
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	cin >> N;
	for (int i = 1; i <= N; i++) cin >> prearr[i];
	for (int i = 1; i < N; i++) {
		int x, y;  cin >> x >> y;
		G[x].emplace_back(y);
		G[y].emplace_back(x);
	}
	dfs1(1, 0);
	dfs2(1, 0, 0, 0, 1);
	init(1, N, 1);

	cin >> Q;
	while (Q--) solve();
}