#469955 (C++14) No.1038 TreeAddQuery

提出ソース

結果

問題	No.1038 TreeAddQuery
ユーザー	QCFium
提出日時	2020-04-24 23:16:45
言語	C++14 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	1,099 ms / 4,000 ms
コード長	4,765 bytes
コンパイル時間	2,673 ms
コンパイル使用メモリ	197,556 KB
実行使用メモリ	46,664 KB
最終ジャッジ日時	2024-05-03 23:15:06
合計ジャッジ時間	16,684 ms
ジャッジサーバーID （参考情報）	judge3 / judge5

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,812 KB
testcase_01	AC	1 ms 6,940 KB
testcase_02	AC	1 ms 6,940 KB
testcase_03	AC	9 ms 6,940 KB
testcase_04	AC	10 ms 6,940 KB
testcase_05	AC	9 ms 6,940 KB
testcase_06	AC	8 ms 6,944 KB
testcase_07	AC	11 ms 6,944 KB
testcase_08	AC	577 ms 27,128 KB
testcase_09	AC	671 ms 26,264 KB
testcase_10	AC	688 ms 30,056 KB
testcase_11	AC	686 ms 30,160 KB
testcase_12	AC	702 ms 32,116 KB
testcase_13	AC	1,076 ms 42,576 KB
testcase_14	AC	979 ms 33,136 KB
testcase_15	AC	926 ms 33,748 KB
testcase_16	AC	936 ms 28,304 KB
testcase_17	AC	876 ms 29,748 KB
testcase_18	AC	175 ms 24,968 KB
testcase_19	AC	218 ms 24,012 KB
testcase_20	AC	193 ms 23,172 KB
testcase_21	AC	231 ms 23,704 KB
testcase_22	AC	284 ms 24,688 KB
testcase_23	AC	319 ms 25,216 KB
testcase_24	AC	488 ms 28,380 KB
testcase_25	AC	1,099 ms 46,664 KB
testcase_26	AC	640 ms 34,644 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include <bits/stdc++.h>

int ri() {
	int n;
	scanf("%d", &n);
	return n;
}
struct Query {
	int range;
	int weight;
	int id;
};

struct AVL {
	struct Node {
		Node *l;
		Node *r;
		int size;
		int height;
		int key;
		int val;
		int64_t sum;
		Node *fetch() {
			size = 1 + l->size + r->size;
			height = 1 + std::max(l->height, r->height);
			sum = val + l->sum + r->sum;
			return this;
		}
		Node *rotate_l() {
			Node *new_root = r;
			r = new_root->l;
			new_root->l = this;
			return fetch(), new_root->fetch();
		}
		Node *rotate_r() {
			Node *new_root = l;
			l = new_root->r;
			new_root->r = this;
			return fetch(), new_root->fetch();
		}
		int height_diff() { return l->height - r->height; }
		Node *balance() {
			int dif = height_diff();
			if (dif == 2) {
				if (l->height_diff() < 0) l = l->rotate_l();
				return rotate_r();
			} else if (dif == -2) {
				if (r->height_diff() > 0) r = r->rotate_r();
				return rotate_l();
			} else return fetch();
		}
	};
	static Node *nodes, *NONE;
	Node *root = NONE;
	static int head;
	AVL () {
		if (!head) nodes[head++] = {NONE, NONE, 0, 0, 0, 0, 0};
	}
	static Node *insert(Node *node, int key, int val) {
		if (node == NONE) return &(nodes[head++] = {NONE, NONE, 1, 1, key, val, val});
		if (key < node->key) node->l = insert(node->l, key, val);
		else if (key > node->key) node->r = insert(node->r, key, val);
		else assert(0);
		return node->balance();
	}
	static int64_t sum(Node *node, int r) {
		if (node == NONE) return 0;
		if (node->key < r) return node->l->sum + node->val + sum(node->r, r);
		else return sum(node->l, r);
	}
	void insert(int key, int val) { root = insert(root, key, val); }
	int64_t sum(int r) { return sum(root, r); }
	static void reset() { if (head) head = 1; }
	void clear() { root = NONE; }
};
AVL::Node *AVL::nodes = new Node[5000000], *AVL::NONE = nodes;
int AVL::head = 0;



std::vector<std::vector<int> > hen;
std::vector<int> size;
std::vector<bool> alive;
std::vector<std::vector<Query> > qs;
void dfs1(int i, int prev) {
	size[i] = qs[i].size() + 1;
	for (auto j : hen[i]) if (j != prev && alive[j]) {
		dfs1(j, i);
		size[i] += size[j];
	}
}
int centroid(int i) {
	dfs1(i, -1);
	int minimum = (size[i] + 1) / 2;
	while (1) {
		int next = -1;
		for (auto j : hen[i]) if (alive[j] && size[j] < size[i] && size[j] >= minimum) next = j;
		if (next == -1) break;
		i = next;
	}
	return i;
}

std::vector<Query> adds, getqs;
std::vector<int> depth;
void dfs2(int i, int prev) {
	for (auto j : qs[i]) {
		if (j.range >= depth[i]) adds.push_back({j.range - depth[i], j.weight, j.id});
		getqs.push_back({depth[i], 0, j.id});
	}
	for (auto j : hen[i]) if (j != prev && alive[j]) {
		depth[j] = depth[i] + 1;
		dfs2(j, i);
	}
}

std::vector<int64_t> res;
void centroid_rec(int i) {
	// std::cerr << i << " " << centroid(i) << std::endl;
	i = centroid(i);
	std::vector<std::pair<int, Query> > add_all, get_all;
	int group = 0;
	for (auto j : hen[i]) if (alive[j]) {
		adds.clear();
		getqs.clear();
		depth[j] = 1;
		dfs2(j, i);
		for (auto i : adds) add_all.push_back({group, i});
		for (auto i : getqs) get_all.push_back({group, i});
		group++;
	}
	for (auto j : qs[i]) {
		add_all.push_back({group, {j.range, j.weight, j.id}});
		get_all.push_back({group, {0, 0, j.id}});
	}
	std::sort(add_all.begin(), add_all.end(), [] (auto &i, auto &j) { return i.second.range > j.second.range; });
	std::sort(get_all.begin(), get_all.end(), [] (auto &i, auto &j) { return i.second.range > j.second.range; });
	/*
	for (auto i : get_all) std::cerr << "  GET " << i.second.range << " " << i.second.weight << " " << i.second.id << std::endl;
	for (auto i : add_all) std::cerr << "  ADD " << i.second.range << " " << i.second.weight << " " << i.second.id << std::endl;
	*/
	int head = 0;
	AVL all;
	AVL each[group + 1];
	for (auto j : get_all) {
		while (head < (int) add_all.size() && add_all[head].second.range >= j.second.range) {
			all.insert(add_all[head].second.id, add_all[head].second.weight);
			each[add_all[head].first].insert(add_all[head].second.id, add_all[head].second.weight);
			head++;
		}
		res[j.second.id] += all.sum(j.second.id);
		if (j.first != group) res[j.second.id] -= each[j.first].sum(j.second.id);
	}
	alive[i] = false;
	for (auto j : hen[i]) if (alive[j]) centroid_rec(j);
	AVL::reset();
}

int main() {
	int n = ri();
	int q = ri();
	hen.resize(n);
	for (int i = 1; i < n; i++) {
		int a = ri() - 1;
		int b = ri() - 1;
		hen[a].push_back(b);
		hen[b].push_back(a);
	}
	qs.resize(n);
	for (int i = 0; i < q; i++) {
		int a = ri() - 1;
		int b = ri();
		int c = ri();
		qs[a].push_back({b, c, i});
	}
	size.resize(n);
	alive.resize(n, true);
	depth.resize(n);
	res.resize(q);
	centroid_rec(0);
	for (auto i : res) printf("%" PRId64 "\n", i);
	return 0;
}

yukicoder

結果

テストケース

ソースコード