P4390 [BalkanOI2007] Mokia 摩基亚

https://www.luogu.com.cn/problem/P4390

摩尔瓦多的移动电话公司摩基亚（Mokia）设计出了一种新的用户定位系统。和其他的定位系统一样，它能够迅速回答任何形如 “用户 C 的位置在哪？” 的问题，精确到毫米。但其真正高科技之处在于，它能够回答形如 “给定区域内有多少名用户？” 的问题。

在定位系统中，世界被认为是一个 $w\times w$ 的正方形区域，由 $1\times 1$ 的方格组成。每个方格都有一个坐标 $(x,y)$，$1\le x,y\le w$。坐标的编号从 $1$ 开始。对于一个 $4\times 4$ 的正方形，就有 $1\le x\le 4$，$1\le y\le 4$（如图）：

请帮助 Mokia 公司编写一个程序来计算在某个矩形区域内有多少名用户。

输入格式

有三种命令，意义如下：

命令	参数	意义
$0$	$w$	初始化一个全零矩阵。本命令仅开始时出现一次。
$1$	$xya$	向方格 $(x,y)$ 中添加 $a$ 个用户。$a$ 是正整数。
$2$	$x_1{y}_1{x}_2{y}_2$	查询 $x_1\le x\le x_2$，$y_1\le y\le y_2$ 所规定的矩形中的用户数量。
$3$	无参数	结束程序。本命令仅结束时出现一次。

输入共若干行，每行有若干个整数，表示一个命令。

输出格式

对所有命令 $2$，输出一个一行整数，即当前询问矩形内的用户数量。

样例 #1

样例输入 #1

0 4
1 2 3 3
2 1 1 3 3
1 2 2 2
2 2 2 3 4
3

样例输出 #1

3
5

提示

数据规模与约定

对于 $100\mathrm{\%}$ 的数据，保证：

• $1\le w\le 2\times {10}^6$。
• $1\le x_1\le x_2\le w$，$1\le y_1\le y_2\le w$，$1\le x,y\le w$，$0<a\le 10000$。
• 命令 $1$ 不超过 $160000$ 个。
• 命令 $2$ 不超过 $10000$ 个。

KD树超时，能AC前两个数据。

class KDTree:
    def __init__(self):
        self.nodes = []
        self.root = None

    class Node:
        def __init__(self, x, y, a):
            self.x = x
            self.y = y
            self.a = a
            self.left = None
            self.right = None
            self.dim = 0  # 分割维度：0 表示 x，1 表示 y

    def build(self, points, depth=0):
        if not points:
            return None
        k = depth % 2  # 当前分割维度
        points.sort(key=lambda p: (p[0], p[1]) if k == 0 else (p[1], p[0]))
        mid = len(points) // 2
        node = self.Node(*points[mid])
        node.dim = k
        node.left = self.build(points[:mid], depth + 1)
        node.right = self.build(points[mid + 1:], depth + 1)
        return node

    def insert(self, root, x, y, a, depth=0):
        if not root:
            return self.Node(x, y, a)
        k = depth % 2
        if (x < root.x if k == 0 else y < root.y):
            root.left = self.insert(root.left, x, y, a, depth + 1)
        else:
            root.right = self.insert(root.right, x, y, a, depth + 1)
        return root

    def query(self, root, x1, y1, x2, y2, depth=0):
        if not root:
            return 0
        k = depth % 2
        if x1 <= root.x <= x2 and y1 <= root.y <= y2:
            res = root.a
        else:
            res = 0
        if k == 0:
            if x1 <= root.x:
                res += self.query(root.left, x1, y1, x2, y2, depth + 1)
            if x2 >= root.x:
                res += self.query(root.right, x1, y1, x2, y2, depth + 1)
        else:
            if y1 <= root.y:
                res += self.query(root.left, x1, y1, x2, y2, depth + 1)
            if y2 >= root.y:
                res += self.query(root.right, x1, y1, x2, y2, depth + 1)
        return res

    def add_point(self, x, y, a):
        self.root = self.insert(self.root, x, y, a)

    def range_query(self, x1, y1, x2, y2):
        return self.query(self.root, x1, y1, x2, y2)


# 主程序
import sys
input = sys.stdin.read
data = input().splitlines()

kdtree = KDTree()

for line in data:
    cmd = list(map(int, line.split()))
    if cmd[0] == 0:
        w = cmd[1]  # 初始化宽度，但不需要存储实际矩阵
        continue
    elif cmd[0] == 1:
        x, y, a = cmd[1:]
        kdtree.add_point(x, y, a)
    elif cmd[0] == 2:
        x1, y1, x2, y2 = cmd[1:]
        print(kdtree.range_query(x1, y1, x2, y2))
    elif cmd[0] == 3:
        break

二维树状数组（Fenwick Tree / BIT）或者二维线段树。能AC第一个数据，其他数据超内存。

import sys

input = sys.stdin.read
data = input().splitlines()

class FenwickTree2D:
    def __init__(self, size):
        self.size = size
        self.tree = [[0] * (size + 1) for _ in range(size + 1)]

    def update(self, x, y, delta):
        while x <= self.size:
            ny = y
            while ny <= self.size:
                self.tree[x][ny] += delta
                ny += ny & -ny
            x += x & -x

    def query(self, x, y):
        sum = 0
        while x > 0:
            ny = y
            while ny > 0:
                sum += self.tree[x][ny]
                ny -= ny & -ny
            x -= x & -x
        return sum

    def range_query(self, x1, y1, x2, y2):
        return (self.query(x2, y2) - self.query(x2, y1 - 1) -
                self.query(x1 - 1, y2) + self.query(x1 - 1, y1 - 1))

def process_commands(data):
    w = 0
    fenwick_tree = None
    for line in data:
        cmd = list(map(int, line.split()))
        if cmd[0] == 0:
            w = cmd[1]
            fenwick_tree = FenwickTree2D(w)
        elif cmd[0] == 1:
            x, y, a = cmd[1:]
            fenwick_tree.update(x, y, a)
        elif cmd[0] == 2:
            x1, y1, x2, y2 = cmd[1:]
            print(fenwick_tree.range_query(x1, y1, x2, y2))
        elif cmd[0] == 3:
            break

process_commands(data)