BZOJ 2989 数列 —— kd-tree + 旋转坐标系

题目链接：点我啊╭(╯^╰)╮

题目大意：

解题思路：

    将 $（i,a[i]）$ 视为一个点，查询即为曼哈顿距离 $\le k$ 的点数
    考虑用 $kd-tree$ 维护，但查询项是一个以 $（i,a[i]）$ 为中心的菱形
    这样单次查询的复杂度无法保证（但是我测了一下只用了500多ms）

    考虑将坐标系旋转 $45°$，菱形就变成了正方形
    只需要将点 $（x,y）$ 变成 $（x-y,x+y）$
    旋转之后，原坐标系的曼哈顿距离就是新坐标系的切比雪夫距离
    这里只需要求一个正方形内的点数即可
    测了一下确实快了很多
    交题链接

不旋转坐标系：

#include<bits/stdc++.h>
#define rint register int
#define deb(x) cerr<<#x<<" = "<<(x)<<'\n';
using namespace std;
typedef long long ll;
typedef pair <int,int> pii;
const int maxn = 2e5 + 5;
int n, q, rt, tot, now, g[maxn], ans;
int sta[maxn], top;
struct node{
	int pla[2];
} a[maxn];
struct tree{
	int mx[2], mn[2], sz, ls, rs;
	node place;
} t[maxn]; 

inline int New(){
	if(top) return sta[top--];
	return ++tot; 
}
inline bool cmp(const node &A, const node &B){
	return A.pla[now] < B.pla[now];
}

inline void update(int p){
	for(int i=0; i<=1; i++){
		t[p].mx[i] = t[p].mn[i] = t[p].place.pla[i];
		if(t[p].ls) t[p].mx[i] = max(t[p].mx[i], t[t[p].ls].mx[i]), \
					t[p].mn[i] = min(t[p].mn[i], t[t[p].ls].mn[i]);
		if(t[p].rs) t[p].mx[i] = max(t[p].mx[i], t[t[p].rs].mx[i]), \
					t[p].mn[i] = min(t[p].mn[i], t[t[p].rs].mn[i]);
	}
	t[p].sz = t[t[p].ls].sz + t[t[p].rs].sz + 1;
}

inline int build(int l, int r, int opt){
	if(l > r) return 0;
	int x = New(), mid = l + r >> 1; now = opt;
	nth_element(a+l, a+mid, a+r+1, cmp); t[x].place = a[mid];
	t[x].ls = build(l, mid-1, opt^1);
	t[x].rs = build(mid+1, r, opt^1);
	update(x); return x;
}

inline void pia(int p, int cnt){
	if(t[p].ls) pia(t[p].ls, cnt);
	a[cnt+t[t[p].ls].sz+1] = t[p].place, sta[++top] = p;
	if(t[p].rs) pia(t[p].rs, cnt+t[t[p].ls].sz+1);
}

inline void check(int &p, int opt){
	if(t[p].sz*0.75<max(t[t[p].ls].sz, t[t[p].rs].sz))
		pia(p, 0), p = build(1, t[p].sz, opt);
}

inline void insert(node ret, int &p, int opt){
	if(!p){
		p = New(); t[p].place = ret;
		t[p].ls = t[p].rs = 0;
		update(p); return;
	}
	if(ret.pla[opt] <= t[p].place.pla[opt]) insert(ret, t[p].ls, opt^1);
	else insert(ret, t[p].rs, opt^1);
	update(p); check(p, opt);
}

inline int ck(node ret, int p, int k){
	int mx = max(abs(ret.pla[0] - t[p].mn[0]), abs(ret.pla[0] - t[p].mx[0])) + \
			 max(abs(ret.pla[1] - t[p].mn[1]), abs(ret.pla[1] - t[p].mx[1]));
	int mn = max(0, ret.pla[0] - t[p].mx[0]) + max(0, t[p].mn[0] - ret.pla[0]) + 
			 max(0, ret.pla[1] - t[p].mx[1]) + max(0, t[p].mn[1] - ret.pla[1]);
	if(mx <= k) return 1;
	if(mn > k) return -1;
	return 0;
}

inline int query(node ret, int p, int k){
	if(!p) return 0;
	int q = ck(ret, p, k);
	if(q == -1) return 0;
	if(q == 1) return t[p].sz;
	int res = 0;
	if(abs(t[p].place.pla[0]-ret.pla[0]) + \
		abs(t[p].place.pla[1]-ret.pla[1]) <= k)
			res++;
	res += query(ret, t[p].ls, k);
	res += query(ret, t[p].rs, k);
	return res;
} 

signed main() {
	scanf("%d%d", &n, &q);
	for(int i=1; i<=n; i++){
		scanf("%d", g+i);
		a[i].pla[0] = i;
		a[i].pla[1] = g[i];
	}
	rt = build(1, n, 0);
	while(q--){
		char s[15]; int x, k; node ins;
		scanf("%s%d%d", s, &x, &k);
		ins.pla[0] = x, ins.pla[1] = k;
		if(s[0] == 'M') g[x] = k, insert(ins, rt, 0);
		else {
			ins.pla[1] = g[x];
			printf("%d\n", query(ins, rt, k));
		}
	}
}

---

旋转坐标系：

#include<bits/stdc++.h>
#define rint register int
#define deb(x) cerr<<#x<<" = "<<(x)<<'\n';
using namespace std;
typedef long long ll;
typedef pair <int,int> pii;
const int maxn = 2e5 + 5;
int n, q, rt, tot, now, g[maxn], ans;
int sta[maxn], top;
struct node{
	int pla[2];
} a[maxn];
struct tree{
	int mx[2], mn[2], sz, ls, rs;
	node place;
} t[maxn]; 

inline int New(){
	if(top) return sta[top--];
	return ++tot; 
}
inline bool cmp(const node &A, const node &B){
	return A.pla[now] < B.pla[now];
}

inline void update(int p){
	for(int i=0; i<=1; i++){
		t[p].mx[i] = t[p].mn[i] = t[p].place.pla[i];
		if(t[p].ls) t[p].mx[i] = max(t[p].mx[i], t[t[p].ls].mx[i]), \
					t[p].mn[i] = min(t[p].mn[i], t[t[p].ls].mn[i]);
		if(t[p].rs) t[p].mx[i] = max(t[p].mx[i], t[t[p].rs].mx[i]), \
					t[p].mn[i] = min(t[p].mn[i], t[t[p].rs].mn[i]);
	}
	t[p].sz = t[t[p].ls].sz + t[t[p].rs].sz + 1;
}

inline int build(int l, int r, int opt){
	if(l > r) return 0;
	int x = New(), mid = l + r >> 1; now = opt;
	nth_element(a+l, a+mid, a+r+1, cmp); t[x].place = a[mid];
	t[x].ls = build(l, mid-1, opt^1);
	t[x].rs = build(mid+1, r, opt^1);
	update(x); return x;
}

inline void pia(int p, int cnt){
	if(t[p].ls) pia(t[p].ls, cnt);
	a[cnt+t[t[p].ls].sz+1] = t[p].place, sta[++top] = p;
	if(t[p].rs) pia(t[p].rs, cnt+t[t[p].ls].sz+1);
}

inline void check(int &p, int opt){
	if(t[p].sz*0.75<max(t[t[p].ls].sz, t[t[p].rs].sz))
		pia(p, 0), p = build(1, t[p].sz, opt);
}

inline void insert(node ret, int &p, int opt){
	if(!p){
		p = New(); t[p].place = ret;
		t[p].ls = t[p].rs = 0;
		update(p); return;
	}
	if(ret.pla[opt] <= t[p].place.pla[opt]) insert(ret, t[p].ls, opt^1);
	else insert(ret, t[p].rs, opt^1);
	update(p); check(p, opt);
}

inline int ck(int p, int x1, int y1, int x2, int y2){
	if(t[p].mn[0]>x2 || t[p].mx[0]<x1) return -1;
	if(t[p].mn[1]>y2 || t[p].mx[1]<y1) return -1;
	if(t[p].mn[0]>=x1 && t[p].mx[0]<=x2 && t[p].mn[1]>=y1 && t[p].mx[1]<=y2) return 1;
	return 0;
}

inline int query(int p, int x1, int y1, int x2, int y2){
	if(!p) return 0;
	int q = ck(p, x1, y1, x2, y2);
	if(q == -1) return 0;
	if(q == 1) return t[p].sz;
	int res = 0;
	if(t[p].place.pla[0]>=x1 && t[p].place.pla[0]<=x2 && \
		 t[p].place.pla[1]>=y1 && t[p].place.pla[1]<=y2) 
			res++;
	res += query(t[p].ls, x1, y1, x2, y2);
	res += query(t[p].rs, x1, y1, x2, y2);
	return res;
} 

signed main() {
	scanf("%d%d", &n, &q);
	for(int i=1; i<=n; i++){
		scanf("%d", g+i);
		a[i].pla[0] = i - g[i];
		a[i].pla[1] = i + g[i];
	}
	rt = build(1, n, 0);
	while(q--){
		char s[15]; int x, k; node ins;
		scanf("%s%d%d", s, &x, &k);
		ins.pla[0] = x - k, ins.pla[1] = x + k;
		if(s[0] == 'M') g[x] = k, insert(ins, rt, 0);
		else printf("%d\n", query(rt, x-g[x]-k, x+g[x]-k, x-g[x]+k, x+g[x]+k));
	}
}