lc3072 将元素分配到两个数组中2
给定数组nums[n]，定义f(arr,val)表示数组arr中大于val的元素个数，需要操作n次将nums分配到两个数组里，具体如下：
- 第1次操作将nums[1]追加到arr1，第2次操作将nums[2]追加到arr2
- 后续第i次操作：
- 如果f(arr1,nums[i]) > f(arr2,nums[i])，则将nums[i]追加到arr1。
- 如果f(arr1,nums[i]) < f(arr2,nums[i])，则将nums[i]追加到arr2。
- 如果相等，则将nums[i]追加到元素少的数组中，如果个数也相等，则追加到arr1。
最后连接arr1和arr2返回。
3<=n<=1e5
一般做法是离散化+树状数组，这里直接套平衡树模板。需要注意的是，力扣在多组数据时的判超时会检查总耗时，因此不要每次分配大的平衡树，好的做法是使用全局的树，每次用时reset。
using ll = long long;
template <typename TYPE>
struct Treap {
struct Node {
TYPE data, sum;
int rnd, siz, dup, son[2];
void init(const TYPE & d) {
data = sum = d;
rnd = rand();
siz = dup = 1;
son[0] = son[1] = 0;
}
};
Treap(size_t sz, bool multi):multiple(multi) {
node.resize(sz);
reset();
}
int newnode(const TYPE & d) {
total += 1;
node[total].init(d);
return total;
}
void reset() { root = total = 0; }
void maintain(int x) {
node[x].siz = node[x].dup;
node[x].sum = node[x].data * node[x].dup;
if (node[x].son[0]) {
node[x].siz += node[node[x].son[0]].siz;
node[x].sum += node[node[x].son[0]].sum;
}
if (node[x].son[1]) {
node[x].siz += node[node[x].son[1]].siz;
node[x].sum += node[node[x].son[1]].sum;
}
}
void rotate(int d, int &r) {
int k = node[r].son[d^1];
node[r].son[d^1] = node[k].son[d];
node[k].son[d] = r;
maintain(r);
maintain(k);
r = k;
}
void insert(const TYPE &data, int &r, bool &ans) {
if (r) {
if (!(data < node[r].data) && !(node[r].data < data)) {
ans = false;
if (multiple) {
node[r].dup += 1;
maintain(r);
}
} else {
int d = data < node[r].data ? 0 : 1;
insert(data, node[r].son[d], ans);
if (node[node[r].son[d]].rnd > node[r].rnd) {
rotate(d^1, r);
} else {
maintain(r);
}
}
} else {
r = newnode(data);
}
}
void getkth(int k, int r, TYPE& data) {
int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
int y = node[r].dup;
if (k <= x) {
getkth(k, node[r].son[0], data);
} else if (k <= x + y) {
data = node[r].data;
} else {
getkth(k-x-y, node[r].son[1], data);
}
}
TYPE getksum(int k, int r) {
if (k <= 0 || r == 0) return 0;
int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
int y = node[r].dup;
if (k <= x) return getksum(k, node[r].son[0]);
if (k <= x+y) return node[node[r].son[0]].sum + node[r].data * (k-x);
return node[node[r].son[0]].sum + node[r].data * y + getksum(k-x-y,node[r].son[1]);
}
void erase(const TYPE& data, int & r) {
if (r == 0) return;
int d = -1;
if (data < node[r].data) {
d = 0;
} else if (node[r].data < data) {
d = 1;
}
if (d == -1) {
node[r].dup -= 1;
if (node[r].dup > 0) {
maintain(r);
} else {
if (node[r].son[0] == 0) {
r = node[r].son[1];
} else if (node[r].son[1] == 0) {
r = node[r].son[0];
} else {
int dd = node[node[r].son[0]].rnd > node[node[r].son[1]].rnd ? 1 : 0;
rotate(dd, r);
erase(data, node[r].son[dd]);
}
}
} else {
erase(data, node[r].son[d]);
}
if (r) maintain(r);
}
int ltcnt(const TYPE& data, int r) {
if (r == 0) return 0;
int x = node[r].son[0] ? node[node[r].son[0]].siz : 0;
if (data < node[r].data) {
return ltcnt(data, node[r].son[0]);
}
if (!(data < node[r].data) && !(node[r].data < data)) {
return x;
}
return x + node[r].dup + ltcnt(data, node[r].son[1]);
}
int gtcnt(const TYPE& data, int r) {
if (r == 0) return 0;
int x = node[r].son[1] ? node[node[r].son[1]].siz : 0;
if (data > node[r].data) {
return gtcnt(data, node[r].son[1]);
}
if (!(data < node[r].data) && !(node[r].data < data)) {
return x;
}
return x + node[r].dup + gtcnt(data, node[r].son[0]);
}
int count(const TYPE& data, int r) {
if (r == 0) return 0;
if (data < node[r].data) return count(data, node[r].son[0]);
if (node[r].data < data) return count(data, node[r].son[1]);
return node[r].dup;
}
void prev(const TYPE& data, int r, TYPE& result, bool& ret) {
if (r) {
if (node[r].data < data) {
if (ret) {
result = max(result, node[r].data);
} else {
result = node[r].data;
ret = true;
}
prev(data, node[r].son[1], result, ret);
} else {
prev(data, node[r].son[0], result, ret);
}
}
}
void next(const TYPE& data, int r, TYPE& result, bool& ret) {
if (r) {
if (data < node[r].data) {
if (ret) {
result = min(result, node[r].data);
} else {
result = node[r].data;
ret = true;
}
next(data, node[r].son[0], result, ret);
} else {
next(data, node[r].son[1], result, ret);
}
}
}
vector<Node> node;
int root, total;
bool multiple;
bool insert(const TYPE& data) {
bool ret = true;
insert(data, root, ret);
return ret;
}
bool kth(int k, TYPE &data) {
if (!root || k <= 0 || k > node[root].siz)
return false;
getkth(k, root, data);
return true;
}
TYPE ksum(int k) {
assert(root && k>0 && k<=node[root].siz);
return getksum(k, root);
}
int count(const TYPE &data) {
return count(data, root);
}
int size() const {
return root ? node[root].siz : 0;
}
void erase(const TYPE& data) {
return erase(data, root);
}
int ltcnt(const TYPE& data) {
return ltcnt(data, root);
}
int gtcnt(const TYPE& data) {
return gtcnt(data, root);
}
int lecnt(const TYPE& data) {
return size() - gtcnt(data, root);
}
int gecnt(const TYPE& data) {
return size() - ltcnt(data, root);
}
bool prev(const TYPE& data, TYPE& result) {
bool ret = false;
prev(data, root, result, ret);
return ret;
}
bool next(const TYPE& data, TYPE& result) {
bool ret = false;
next(data, root, result, ret);
return ret;
}
};
const int N = 100005;
Treap<ll> t1(N,1), t2(N,1);
class Solution {
public:
vector<int> resultArray(vector<int>& nums) {
int n = nums.size();
vector<int> z1(n), z2(n);
int n1 = 0, n2 = 0;
t1.reset();
t2.reset();
z1[n1++] = nums[0]; t1.insert(nums[0]);
z2[n2++] = nums[1]; t2.insert(nums[1]);
for (int i = 2; i < n; i++) {
int v1 = t1.gtcnt(nums[i]);
int v2 = t2.gtcnt(nums[i]);
int to = 0;
if (v1 > v2) {
to = 1;
} else if (v1 < v2) {
to = 2;
} else {
if (n1 <= n2) {
to = 1;
} else {
to = 2;
}
}
if (to == 1) {
z1[n1++] = nums[i]; t1.insert(nums[i]);
} else {
z2[n2++] = nums[i]; t2.insert(nums[i]);
}
}
for (int i = 0; i < n2; i++) {
z1[n1++] = z2[i];
}
return z1;
}
};
