是否有任何数据结构能够返回数组的前缀和[1]、更新元素以及向数组中插入/删除元素，所有这些都在 O(log n) 内？
[1]“前缀总和”是从第一个元素到给定索引的所有元素的总和
例如，给定非负整数数组8 1 10 7，前三个元素的前缀和为19 (8 + 1 + 10)。将第一个元素更新为7，插入3 作为第二个元素并删除第三个元素，得到7 3 10 7。同样，前三个元素的前缀总和将是 20。
对于前缀求和和更新，有Fenwick tree。但我不知道如何用它处理 O(log n) 中的添加/删除。
另一方面，有几个二叉搜索树，例如Red-black tree，它们都以对数时间处理更新/插入/删除。但我不知道如何维护给定的顺序并在 O(log n) 中进行前缀求和。
【问题讨论】：
标签： algorithm data-structures binary-search-tree prefix-sum
带有隐式键的treap 可以在每次查询的O(log n) 时间内执行所有这些操作。隐式键的想法非常简单：我们不在节点中存储任何键。相反，我们维护所有节点的子树大小，并在使用此信息添加或删除元素时找到适当的位置。
这是我的实现：
#include <iostream>
#include <memory>
struct Node {
int priority;
int val;
long long sum;
int size;
std::shared_ptr<Node> left;
std::shared_ptr<Node> right;
Node(long val):
priority(rand()), val(val), sum(val), size(1), left(), right() {}
};
// Returns the size of a node owned by t if it is not empty and 0 otherwise.
int getSize(std::shared_ptr<Node> t) {
if (!t)
return 0;
return t->size;
}
// Returns the sum of a node owned by t if it is not empty and 0 otherwise.
long long getSum(std::shared_ptr<Node> t) {
if (!t)
return 0;
return t->sum;
}
// Updates a node owned by t if it is not empty.
void update(std::shared_ptr<Node> t) {
if (t) {
t->size = 1 + getSize(t->left) + getSize(t->right);
t->sum = t->val + getSum(t->left) + getSum(t->right);
}
}
// Merges the nodes owned by L and R and returns the result.
std::shared_ptr<Node> merge(std::shared_ptr<Node> L,
std::shared_ptr<Node> R) {
if (!L || !R)
return L ? L : R;
if (L->priority > R->priority) {
L->right = merge(L->right, R);
update(L);
return L;
} else {
R->left = merge(L, R->left);
update(R);
return R;
}
}
// Splits a subtree rooted in t by pos.
std::pair<std::shared_ptr<Node>, std::shared_ptr<Node>> split(
std::shared_ptr<Node> t,
int pos, int add) {
if (!t)
return make_pair(std::shared_ptr<Node>(), std::shared_ptr<Node>());
int cur = getSize(t->left) + add;
std::pair<std::shared_ptr<Node>, std::shared_ptr<Node>> res;
if (pos <= cur) {
auto ret = split(t->left, pos, add);
t->left = ret.second;
res = make_pair(ret.first, t);
} else {
auto ret = split(t->right, pos, cur + 1);
t->right = ret.first;
res = make_pair(t, ret.second);
}
update(t);
return res;
}
// Returns a prefix sum of [0 ... pos]
long long getPrefixSum(std::shared_ptr<Node>& root, int pos) {
auto parts = split(root, pos + 1, 0);
long long res = getSum(parts.first);
root = merge(parts.first, parts.second);
return res;
}
// Adds a new element at a position pos with a value newValue.
// Indices are zero-based.
void addElement(std::shared_ptr<Node>& root, int pos, int newValue) {
auto parts = split(root, pos, 0);
std::shared_ptr<Node> newNode = std::make_shared<Node>(newValue);
auto temp = merge(parts.first, newNode);
root = merge(temp, parts.second);
}
// Removes an element at the given position pos.
// Indices are zero-based.
void removeElement(std::shared_ptr<Node>& root, int pos) {
auto parts1 = split(root, pos, 0);
auto parts2 = split(parts1.second, 1, 0);
root = merge(parts1.first, parts2.second);
}
int main() {
std::shared_ptr<Node> root;
int n;
std::cin >> n;
for (int i = 0; i < n; i++) {
std::string s;
std::cin >> s;
if (s == "add") {
int pos, val;
std::cin >> pos >> val;
addElement(root, pos, val);
} else if (s == "remove") {
int pos;
std::cin >> pos;
removeElement(root, pos);
} else {
int pos;
std::cin >> pos;
std::cout << getPrefixSum(root, pos) << std::endl;
}
}
return 0;
}
【讨论】：
一个想法：修改 AVL 树。添加和删除是通过索引完成的。每个节点都保存每个子树的计数和总和，以允许 O(log n) 中的所有操作。
实施了 add_node 和 update_node 和 prefix_sum 的概念验证：
class Node:
def __init__(self, value):
self.value = value
self.left = None
self.right = None
self.left_height = 0
self.right_height = 0
self.left_count = 1
self.left_sum = value
self.right_count = 0
self.right_sum = 0
def set_value(self, value):
self.value = value
self.left_sum = self.left.left_sum + self.left.right_sum+self.value if self.left else self.value
def set_left(self, node):
self.left = node
self.left_height = max(node.left_height, node.right_height)+1 if node else 0
self.left_count = node.left_count + node.right_count+1 if node else 1
self.left_sum = node.left_sum + node.right_sum+self.value if node else self.value
def set_right(self, node):
self.right = node
self.right_height = max(node.left_height, node.right_height)+1 if node else 0
self.right_count = node.left_count + node.right_count if node else 0
self.right_sum = node.left_sum + node.right_sum if node else 0
def rotate_left(self):
b = self.right
self.set_right(b.left)
b.set_left(self)
return b
def rotate_right(self):
a = self.left
self.set_left(a.right)
a.set_right(self)
return a
def factor(self):
return self.right_height - self.left_height
def add_node(root, index, node):
if root is None: return node
if index < root.left_count:
root.set_left(add_node(root.left, index, node))
if root.factor() < -1:
if root.left.factor() > 0:
root.set_left(root.left.rotate_left())
return root.rotate_right()
else:
root.set_right(add_node(root.right, index-root.left_count, node))
if root.factor() > 1:
if root.right.factor() < 0:
root.set_right(root.right.rotate_right())
return root.rotate_left()
return root
def update_node(root, index, value):
if root is None: return root
if index+1 < root.left_count:
root.set_left(update_node(root.left, index, value))
elif index+1 > root.left_count:
root.set_right(update_node(root.right, index - root.left_count, value))
else:
root.set_value(value)
return root
def prefix_sum(root, index):
if root is None: return 0
if index+1 < root.left_count:
return prefix_sum(root.left, index)
else:
return root.left_sum + prefix_sum(root.right, index-root.left_count)
import random
tree = None
tree = add_node(tree, 0, Node(10))
tree = add_node(tree, 1, Node(40))
tree = add_node(tree, 1, Node(20))
tree = add_node(tree, 2, Node(70))
tree = update_node(tree, 2, 30)
print prefix_sum(tree, 0)
print prefix_sum(tree, 1)
print prefix_sum(tree, 2)
print prefix_sum(tree, 3)
print prefix_sum(tree, 4)
【讨论】：