Files
DataStructureAndAlgorithm/2023206/main.cpp
2023-11-29 21:08:28 +08:00

140 lines
3.4 KiB
C++

#include <stdio.h>
#define TREE_SIZE 65536
#define NOTHING 65535
struct TreeNode {
// int val;
unsigned short size;
unsigned short left;
unsigned short right;
};
TreeNode tree[TREE_SIZE] = {0};
int vals[TREE_SIZE] = {0};
int tree_root = NOTHING;
int M, K, H;
long long total = 0;
int l = 0;
int insert_node(int num, int pos) {
if (pos == NOTHING) {
// The tree_root is -1!
tree[l % K] = TreeNode{1, NOTHING, NOTHING};
tree_root = l % K;
vals[l % K] = num;
return l % K;
}
tree[pos].size++;
if (vals[pos] > num) {
if (tree[pos].left == NOTHING) {
tree[pos].left = l % K;
tree[l % K] = TreeNode{1, NOTHING, NOTHING};
vals[l % K] = num;
return l % K;
}
else {
return insert_node(num, tree[pos].left);
}
}
else {
if (tree[pos].right == NOTHING) {
tree[pos].right = l % K;
tree[l % K] = TreeNode{1, NOTHING, NOTHING};
vals[l % K] = num;
return l % K;
}
else {
return insert_node(num, tree[pos].right);
}
}
}
int delete_node(int num, int pos) {
if (pos == NOTHING) {
return NOTHING;
}
if (vals[pos] < num) {
tree[pos].size--;
tree[pos].right = delete_node(num, tree[pos].right);
return pos;
}
if (vals[pos] > num) {
tree[pos].size--;
tree[pos].left = delete_node(num, tree[pos].left);
return pos;
}
// If reached here, tree[pos].val == num, so we are going to delete tree[pos]
if (tree[pos].left == NOTHING) {
// The case when left is null, right is something
// and also the case when both left and right is null
return tree[pos].right;
}
if (tree[pos].right == NOTHING) {
// The case when left is something while right is null
return tree[pos].left;
}
// Here, we deal with the case that both children exist.
int parent = pos;
int child = tree[parent].right;
tree[parent].size--;
while (tree[child].left < NOTHING) {
parent = child;
child = tree[child].left;
tree[parent].size--;
}
if (parent == pos) {
tree[parent].right = tree[child].right;
}
else {
tree[parent].left = tree[child].right;
}
tree[child].left = tree[pos].left;
tree[child].right = tree[pos].right;
tree[child].size = tree[pos].size;
return child;
}
int count_less_than(long long target, int pos) {
if (pos == NOTHING) {
return 0;
}
if (vals[pos] < target) {
if (tree[pos].left == NOTHING) {
return 1 + count_less_than(target, tree[pos].right);
}
return tree[tree[pos].left].size + 1 + count_less_than(target, tree[pos].right);
}
if (vals[pos] > target) {
return count_less_than(target, tree[pos].left);
}
if (tree[pos].left == NOTHING) {
return 1;
}
return 1 + tree[tree[pos].left].size;
}
int main() {
scanf("%d %d %d", &M, &K, &H);
int num = 0;
for (l = 0; l < M; l++) {
scanf("%d", &num);
total += count_less_than(((long long)num) + H, tree_root) -
count_less_than(num - H - 1, tree_root);
if (l >= K) {
tree_root = delete_node(vals[l % K], tree_root);
}
insert_node(num, tree_root);
}
printf("%lld\n", total);
return 0;
}