数据中心能耗分析｜进阶线段树

T8 - 数据中心能耗分析

题目链接跳转：点击跳转

本文仅针对对线段树有一定了解且并且有着较高的程序设计能力的选手。本文的前置知识如下：

1. 线段树的构造与维护 - 可以参考文章 浅入线段树与区间最值问题。
2. 初中数学 - 完全平方和公式和完全立方和公式。
3. 取模 - 之如何保证对所有整数取模后的结果为非负整数 - 可以参考本题的 说明/提示 部分。

原本出题的时候我并不知道这道题在许多 OJ 上是有原题的，so sad（下次再改进）。

题目本身应该非常好理解，就是给定一个数组，让你设计一个程序，对程序进行区间求立方和和区间修改的操作。但本题的数据量比较大，$N, M$ 最大可以达到 $10^5$，对于每一次修改和查询都是 $O(N)$ 的时间复杂度，显然用暴力的方法时间复杂度绝对会超时，最高会到 $O(N^2 * M)$ （大概需要 $115$ 天的时间才可以搞定一个测试点）。当看到区间查询和维护操作的时候，不难想到用线段树的方法，在线段树中，单次操作的时间复杂度约为 $O(log_2 N)$，即使当 $N$ 非常大的时候线段树也可以跑得飞起。

解题思路

不得不说，这是一道比较恶心的线段树区间维护的题目。不光写起来比较费劲，而且维护操作运算量比较大。稍有不慎就会写歪（因此写这道题的时候要集中注意力，稍微一个不起眼的问题就容易爆 $0$）。

本题的主要难点就是对一个区间内进行批量区间累加的操作。很容易就想到跟完全立方公式的联系：$(a+b)^3 = a^3 +3a^2b + 3ab^2 + b^3$。区间累加操作也只不过是对区间的所有数字都进行该操作，并对所有操作的结果求和就是该区间进行操作后的立方和。化简可得：

$$\begin{align} \mathtt{ANS} &= \sum_{i=1}^{n}{(a_i+x)^3}\\ &= (a_1+b)^3+(a_2+b)^3+ \cdots + (a_n+b)^3 \\ &= (a_1^3 + 3a_1^2b + 3a_1b^2 + b^3) + (a_2^3 + 3a_2^2b + 3a_2b^2 + b^3) + \cdots + (a_n^3 + 3a_n^2b + 3a_nb^2 + b^3) \\ &= (a_1^3 + a_2^3 + \cdots + a_n^3) + 3b(a_1^2 + a_2^2 + \cdots + a_n^2) + 3b^2(a_1 + a_2 + \cdots + a_n) + nb^3 \\ &= \sum_{i=1}^{n}{a_i^3} + 3b\sum_{i=1}^{n}{a_i^2} + 3b^2\sum_{i=1}^{n}{a_i} + nb^3 \end{align}$$

综上所述，我们只需要用线段树维护三个字段，分别是区间的立方和、区间的平方和以及区间和。在维护平方和的过程中与维护立方和类似，根据完全平方公式 $(a+b)^2 = a^2 + 2ab + b^2$。经过累加和化简可得：

$$\begin{align} \mathtt{ANS} &= \sum_{i=1}^{n}{(a_i+x)^2}\\ &= (a_1+b)^2 + (a_2+b)^2 + \cdots + (a_n+b)^2 \\ &= (a_1^2 + 2a_1b + b^2) + (a_2^2 + 2a_2b + b^2) + \cdots + (a_n^2 + 2a_nb + b^2) \\ &= (a_1^2 + a_2^2 + \cdots + a_n^2) + 2b(a_1 + a_2 + \cdots + a_n) + nb^2 \\ &= \sum_{i=1}^{n}{a_i^2} + 2b\sum_{i=1}^{n}{a_i} + nb^2 \end{align}$$

以上三个字段可以在构造线段树的时候一并初始化，之后的每次更新直接修改懒标记就可以了。一切都交给 push_down() 函数。在每次区间查询和修改之前都进行懒标记下放操作，对区间进行维护。具体维护操作如下：

// rt 是父节点，l和r是rt的两个子节点，len是rt节点区间的长度。
// 其中，(len - len / 2)是l区间的长度，(len / 2)是r区间的长度。
void push_down(Node &rt, Node &l, Node &r, int len){
    if (rt.tag){
        int num = rt.tag;
		// 维护立方和
        l.s3 += 3 * num * l.s2 + 3 * num * num * l.s1 + (len - len / 2) * num * num * num;
        r.s3 += 3 * num * r.s2 + 3 * num * num * r.s1 + (len / 2) * num * num * num;
		
        //维护平方和
        l.s2 += 2 * num * l.s1 + (len - len / 2) * num * num;
        l.s2 += 2 * num * r.s1 + (len / 2) * num * num;

        // 维护区间总和
        l.s1 += (len - len / 2) * num;
        r.s1 += (len / 2) * num;
		
        // 将标记下放到两个子区间
        l.tag += num;
        r.tag += num;
        rt.tag = 0;  // 清空标记。
    }
    return ;
}

注意事项

1. 请注意取模，为了保证答案正确性，请在每一步操作的时候都对结果取模。
2. 开 long long，不然的话只能过前三个测试点（出题人还是挺好的，留了三个小的测试点骗粉）。
3. 在维护立方和、平方和以及和的时候，请注意维护的顺序。应当先维护立方和，再维护平方和，最后再维护区间总和。
4. 注意线段树数组的大小，应当为 $4 \times N$。
5. 建议使用读入优化，直接使用 cin 的效率比 std 慢约 $100\%$​。

时间复杂度

线段树单次查询和修改的复杂度约为 $O(log_2 N)$，初始化的时间复杂度为 $\Theta(N)$，因此本代码的整体时间复杂度可以用多项式 $\Theta(N + M \cdot log_2(N))$ 来表示，整体代码的时间复杂度就为 $O(M \cdot log_2(N))$。在极限数据下，程序只需要 $160ms$ 就可以完成暴力一整年所需的工作。

代码实现

1. 代码使用了宏定义，方便后期进行调式。
2. 以下代码与普通的线段树无太大区别，但请着重关注 push_down() 下放操作。

#include <iostream>
#include <algorithm>
using namespace std;

const int N = 1e5 + 5;
const int MOD = 1e9 + 7;
// 宏定义：lc和rc分别表示左儿子和右儿子在数组中的索引。
#define lc root << 1
#define rc root << 1 | 1
#define int long long
int n, m, k, x, y, v;

struct Node {
    // 分别表示总和、平方和与立方和。
    int s1, s2, s3;
    int tag;
} tree[N << 2];

// 合并区间，直接将两个子区间相加即可。
void push_up(int root) {
    tree[root].s1 = (tree[lc].s1 + tree[rc].s1) % MOD;
    tree[root].s2 = (tree[lc].s2 + tree[rc].s2) % MOD;
    tree[root].s3 = (tree[lc].s3 + tree[rc].s3) % MOD;
    return;
}

// 下放操作，确实码量比较大。在借鉴的时候需要仔细点。
void push_down(Node &rt, Node &l, Node &r, int len) {
    if (rt.tag) {
        int num = rt.tag;

        l.s3 = (l.s3 + 3 * num % MOD * l.s2 % MOD + 3 * num % MOD * num % MOD * l.s1 % MOD + (len - len / 2) * num % MOD * num % MOD * num % MOD) % MOD;
        l.s3 = (l.s3 + MOD) % MOD;
        r.s3 = (r.s3 + 3 * num % MOD * r.s2 % MOD + 3 * num % MOD * num % MOD * r.s1 % MOD + (len / 2) * num % MOD * num % MOD * num % MOD) % MOD;
        r.s3 = (r.s3 + MOD) % MOD;

        l.s2 = (l.s2 + 2 * num % MOD * l.s1 % MOD + (len - len / 2) * num % MOD * num % MOD) % MOD;
        l.s2 = (l.s2 + MOD) % MOD;
        r.s2 = (r.s2 + 2 * num % MOD * r.s1 % MOD + (len / 2) * num % MOD * num % MOD) % MOD;
        r.s2 = (r.s2 + MOD) % MOD;

        l.s1 = (l.s1 + (len - len / 2) * num % MOD) % MOD;
        l.s1 = (l.s1 + MOD) % MOD;
        r.s1 = (r.s1 + (len / 2) * num % MOD) % MOD;
        r.s1 = (r.s1 + MOD) % MOD;

        l.tag = (l.tag + num) % MOD;
        l.tag = (l.tag + MOD) % MOD;
        r.tag = (r.tag + num) % MOD;
        r.tag = (r.tag + MOD) % MOD;
        rt.tag = 0;
    }
    return;
}

// 非常简单的造树操作。
void build(int l, int r, int root) {
    if (l == r) {
        int t; cin >> t;
        tree[root].s1 = t % MOD; 
        tree[root].s2 = t * t % MOD;
        tree[root].s3 = t * t % MOD * t % MOD;
        return;
    }
    int mid = (l + r) >> 1;
    build(l, mid, lc);
    build(mid + 1, r, rc);
    push_up(root);
    return;
}

// 更新操作。
void update(int l, int r, int v, int L, int R, int root) {
    // 跟push_down()函数基本类似。
    if (L <= l && r <= R) {
        tree[root].tag = (tree[root].tag + v) % MOD;
        tree[root].tag = (tree[root].tag + MOD) % MOD;
        tree[root].s3 = (tree[root].s3 + 3 * v % MOD * tree[root].s2 % MOD + 3 * v % MOD * v % MOD * tree[root].s1 % MOD + (r - l + 1) * v % MOD * v % MOD * v % MOD) % MOD;
        tree[root].s3 = (tree[root].s3 + MOD) % MOD;
        tree[root].s2 = (tree[root].s2 + 2 * v % MOD * tree[root].s1 % MOD + (r - l + 1) * v % MOD * v % MOD) % MOD;
        tree[root].s2 = (tree[root].s2 + MOD) % MOD;
        tree[root].s1 = (tree[root].s1 + (r - l + 1) * v % MOD) % MOD;
        tree[root].s1 = (tree[root].s1 + MOD) % MOD;
        return;
    }
    push_down(tree[root], tree[lc], tree[rc], r - l + 1);
    int mid = (l + r) >> 1;
    if (L <= mid) update(l, mid, v, L, R, lc);
    if (R >= mid + 1) update(mid + 1, r, v, L, R, rc);
    push_up(root);
}

// 区间查询操作。
int query(int l, int r, int L, int R, int root) {
    if (L <= l && r <= R)
        return (tree[root].s3 + MOD) % MOD;
    int sum = 0;
    push_down(tree[root], tree[lc], tree[rc], r - l + 1);
    int mid = (l + r) >> 1;
    if (L <= mid) sum = (sum + query(l, mid, L, R, lc)) % MOD;
    if (R >= mid + 1) sum = (sum + query(mid + 1, r, L, R, rc)) % MOD;
    return (sum + MOD) % MOD;
}

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    cin >> n >> m;
    build(1, n, 1);
    while (m--) {
        cin >> k >> x >> y;
        if (k == 1) {
            cin >> v;
            update(1, n, v, x, y, 1);
        } else cout << query(1, n, x, y, 1) << endl;
    }
    return 0;
}

以下是本题的 Python 代码，但由于 Python 常数过大，没有办法通过所有的测试点：

MOD = 10**9 + 7
N = int(1e5 + 5)

class Node:
    def __init__(self):
        self.s1 = 0
        self.s2 = 0
        self.s3 = 0
        self.tag = 0

tree = [Node() for _ in range(N * 4)]

def push_up(root):
    tree[root].s1 = (tree[root * 2].s1 + tree[root * 2 + 1].s1) % MOD
    tree[root].s2 = (tree[root * 2].s2 + tree[root * 2 + 1].s2) % MOD
    tree[root].s3 = (tree[root * 2].s3 + tree[root * 2 + 1].s3) % MOD

def push_down(root, l, r, length):
    if tree[root].tag != 0:
        num = tree[root].tag % MOD
        left_child = root * 2
        right_child = root * 2 + 1
        left_len = length - length // 2
        right_len = length // 2

        # Left child updates
        tree[left_child].s3 = (tree[left_child].s3 + (3 * num * tree[left_child].s2 % MOD) + (3 * num * num % MOD * tree[left_child].s1 % MOD) + (left_len * num % MOD * num % MOD * num % MOD)) % MOD
        tree[left_child].s2 = (tree[left_child].s2 + (2 * num * tree[left_child].s1 % MOD) + (left_len * num % MOD * num % MOD)) % MOD
        tree[left_child].s1 = (tree[left_child].s1 + (left_len * num % MOD)) % MOD
        tree[left_child].tag = (tree[left_child].tag + num) % MOD

        # Right child updates
        tree[right_child].s3 = (tree[right_child].s3 + (3 * num * tree[right_child].s2 % MOD) + (3 * num * num % MOD * tree[right_child].s1 % MOD) + (right_len * num % MOD * num % MOD * num % MOD)) % MOD
        tree[right_child].s2 = (tree[right_child].s2 + (2 * num * tree[right_child].s1 % MOD) + (right_len * num % MOD * num % MOD)) % MOD
        tree[right_child].s1 = (tree[right_child].s1 + (right_len * num % MOD)) % MOD
        tree[right_child].tag = (tree[right_child].tag + num) % MOD

        tree[root].tag = 0

def build(l, r, root):
    if l == r:
        t = data[l - 1] % MOD
        tree[root].s1 = t
        tree[root].s2 = t * t % MOD
        tree[root].s3 = t * t % MOD * t % MOD
        return
    mid = (l + r) // 2
    build(l, mid, root * 2)
    build(mid + 1, r, root * 2 + 1)
    push_up(root)

def update(l, r, v, L, R, root):
    if L <= l and r <= R:
        num = v % MOD
        length = r - l + 1
        tree[root].tag = (tree[root].tag + num) % MOD
        tree[root].s3 = (tree[root].s3 + (3 * num * tree[root].s2 % MOD) + (3 * num * num % MOD * tree[root].s1 % MOD) + (length * num % MOD * num % MOD * num % MOD)) % MOD
        tree[root].s2 = (tree[root].s2 + (2 * num * tree[root].s1 % MOD) + (length * num % MOD * num % MOD)) % MOD
        tree[root].s1 = (tree[root].s1 + (length * num % MOD)) % MOD
        return
    push_down(root, l, r, r - l + 1)
    mid = (l + r) // 2
    if L <= mid:
        update(l, mid, v, L, R, root * 2)
    if R > mid:
        update(mid + 1, r, v, L, R, root * 2 + 1)
    push_up(root)

def query(l, r, L, R, root):
    if L <= l and r <= R:
        return tree[root].s3 % MOD
    push_down(root, l, r, r - l + 1)
    mid = (l + r) // 2
    res = 0
    if L <= mid:
        res = (res + query(l, mid, L, R, root * 2)) % MOD
    if R > mid:
        res = (res + query(mid + 1, r, L, R, root * 2 + 1)) % MOD
    return res % MOD

if __name__ == '__main__':
    import sys
    sys.setrecursionlimit(1 << 25)
    n, m = map(int, sys.stdin.readline().split())
    data = list(map(int, sys.stdin.readline().split()))
    build(1, n, 1)
    for _ in range(m):
        tmp = sys.stdin.readline().split()
        if not tmp:
            continue
        k = int(tmp[0])
        x = int(tmp[1])
        y = int(tmp[2])
        if k == 1:
            v = int(tmp[3])
            update(1, n, v, x, y, 1)
        else:
            print(query(1, n, x, y, 1))