Atcoder解いてみた AtCoder Beginner Contest D - Swap and Range Sum

AtCoder Beginner Contest D - Swap and Range Sum

問題はこちら

回答

素直にやろうとすると、各クエリに対して、入れ替え処理もしくは合計値算出を行うため、計算量がO(Q N)となりアウト
クエリのループは必要なため、O(Q logN)には計算量を抑えたい ※O(Q)では多分解けない
特定範囲内の合計処理なため、セグメントツリーを疑った
合計値算出は(logN)でできる
入れ替え処理も、今回はAx, Ax+1の入れ替えなため、2回木の更新を行えばいいので、こちらもO(logN)でできる
全体の計算量もO(Q logN)なため、これでOK

解説ではセグ木を使わず、O(Q+N)で解いていたので、次回はそれで解く

#include <iostream>
#include <string>
#include <map>
#include <unordered_map>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <vector>
#include <queue>
#include <stack>
#include <limits.h>
#include <bitset>
#include <list>
#include <set>
#include <numeric>
#include <tuple>

int N; // 2~2x10^5 数列の長さ
int Q; // 1~10^5 クエリの個数
std::vector<long long> A; // 整数列 各値は、1~10^4

class SegTree {
private:
	int leafIdxOffset;
	long long defVal;
	std::vector<long long> tree;

public:
	SegTree(const std::vector<long long>& numArr, long long defVal) :defVal(defVal) {
		// 木の要素数を算出
		int numTreeElem = 1;

		while (numTreeElem < numArr.size()) {
			numTreeElem *= 2;
		}

		numTreeElem *= 2;
		tree.resize(numTreeElem);

		leafIdxOffset = numTreeElem / 2;

		// 木の葉に数値をセット
		for (int i = 0; i < numArr.size(); i++) {
			tree[i + leafIdxOffset] = numArr[i]; // treeへのアクセスは1-indexed
		}

		// 葉からノードの値を求める
		for (int i = leafIdxOffset - 1; i >= 1; i--) {
			// 子の座標
			int sonLeft = i * 2;
			int sonRight = i * 2 + 1;

			tree[i] = tree[sonLeft] + tree[sonRight];
		}
	}

	// idxはオフセットを考慮しない葉へのインデックスで、0-indexedで指定する
	void UpdateTree(int idx) {
		idx += leafIdxOffset;

		// 親をたどっていき、更新
		while (idx > 1) {
			idx /= 2; // 親の座標
			int sonLeft = idx * 2;
			int sonRight = idx * 2 + 1;

			tree[idx] = tree[sonLeft] + tree[sonRight];
		}
	}

	void SwapTree(int idx1, int idx2) {
		int treeIdx1 = idx1 + leafIdxOffset;
		int treeIdx2 = idx2 + leafIdxOffset;

		long long valIdx1 = tree[treeIdx1];
		long long valIdx2 = tree[treeIdx2];

		tree[treeIdx1] = valIdx2;
		UpdateTree(idx1);

		tree[treeIdx2] = valIdx1;
		UpdateTree(idx2);
	}

	//[a,b)
	long long Query(int a, int b) {
		long long vLeft = defVal; // 木の根から見たときに左側の合計値
		long long vRight = defVal; // 木の根から見たときに右側の合計値

		for (int left = a + leafIdxOffset, right = b + leafIdxOffset; left < right; left >>= 1, right >>= 1) {
			// クエリの[左,右)が奇数だった場合、余計な値を合計値に加えないための対応
			if (left & 1) vLeft = vLeft + tree[left++];
			if (right & 1) vRight = vRight + tree[--right];
		}
		
		return vLeft + vRight;
	}
};

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

	// クエリ
	// 1 x : Ax, Ax+1の入れ替え
	// 2 l r : l<=x<=rの合計値

	std::cin >> N >> Q;

	A.resize(N);
	for (int i = 0; i < N; i++) {
		std::cin >> A[i];
	}

	SegTree segTree(A, 0LL);

	std::vector<long long> ans;
	for (int q = 0; q < Q; q++) {
		int com;
		std::cin >> com;

		int x, l, r;

		if (com == 1) {
			std::cin>>x;
			x--;

			// xと、x+1を入れ替え
			segTree.SwapTree(x, x + 1);
		}
		else {
			std::cin >> l;
			std::cin >> r;
			l--;
			r--;

			long long val = segTree.Query(l, r + 1); // 半開区間なため、r+1
			ans.push_back(val);
		}
	}

	for (int i = 0; i < ans.size(); i++) {
		std::cout << ans[i] << std::endl;
	}

	return 0;
}