tags: DSA C++ LeetCode

写在前面

总结一下经典的差分数组方法, (华为机试刚考了), 思路很简单, 但是没遇到的话想写出来还是有点难度的.

参考了 labuladong 的博客, 里面的代码是 Java 实现的, 这里用 C++重写一下.

小而美的算法技巧：差分数组.

思路

差分数组是一种支持频繁对数组的某一区间进行增减, 以及查询的数据结构.

基本思想就是: (高中学的数列)

设数列为 { a i } , i ∈ [ 0 , n ] \\{a_i\\},\\ i\\in[0,\\ n] {ai​}, i∈[0, n], 那么做一新数列 { b j } , j ∈ [ 0 , n ] \\{b_j\\},\\ j\\in[0, n] {bj​}, j∈[0,n], 且满足
{ b 0 = a 0 b j = a i − a i − 1 , j ∈ [ 1 , n ] \\begin{cases} b_0 = a_0\\\\ b_j = a_i - a_{i-1}, \\ j\\in[1,\\ n] \\end{cases} {b0​=a0​bj​=ai​−ai−1​, j∈[1, n]​
这样的数列(数组) b j b_j bj​ 其实就是差分数组.

那么为什么这样的数组支持对区间的频繁增减呢?

例子

A = { 4 , 7 , 5 , 2 , 9 } A=\\{4,7,5,2,9\\} A={4,7,5,2,9}, 对应的差分数组 B = { 4 , 3 , − 2 , − 3 , 7 } B=\\{4,3,-2,-3,7\\} B={4,3,−2,−3,7}.

初始化

void init(vector<int> &A, vector<int> & B) {B[0] = A[0];for (int i{1}; i < A.size(); ++i) B[i] = A[i] - A[i - 1];
}

那么有了差分数组, 计算区间和就像前缀和一样简洁了:

查询: 返回原始数组

直接将上述操作反过来即可:

vector<int> get_raw(vector<int> &B) {int n = B.size();vector<int> C(B);for (int i{1}; i < B.size(); ++i) C[i] = B[i] + C[i - 1];return C;
}

修改: 增减

如果要对区间(闭区间) $[2,3]$ 每一个数都增加 1, 那么只需要修改区间的第一个下标位置的元素(加一)和最后一个下标位置的下一个元素(减一)即可, 即
B ′ = { 4 , 3 , − 1 , − 3 , 6 } B' = \\{4,3,-1,-3,6\\} B′={4,3,−1,−3,6}

一般的, 对差分数组$B$的区间$[i,j]$进行增(减)$k$的操作, 只需将$B[i]\pm =k$且$B[j+1]\mp =k$即可.

void modify(vector<int> &B, int i, int j, int k) {B[i] += k;if (j + 1 < B.size()) // 考虑数组越界B[j + 1] -= k;
}

封装成类

头文件与打印重载:

#include <bits/stdc++.h>
using namespace std;ostream& operator<<(ostream& os, const vector<int>& v) {for (auto i : v) os << i << " ";return os << endl;
}

类: (声明与定义写在一起了)

template <typename T>
class DiffArray {vector<T> arr;public:DiffArray() : arr() {}DiffArray(vector<T> _arr) { init(_arr); }void init(const vector<T>& raw_arr) {int n = raw_arr.size();arr.resize(n);arr[0] = raw_arr[0];for (int i{1}; i < n; ++i) arr[i] = raw_arr[i] - raw_arr[i - 1];}vector<T> get_raw() {vector<T> raw_arr(arr);for (int i{1}; i < arr.size(); ++i)raw_arr[i] = arr[i] + raw_arr[i - 1];return raw_arr;}void modify(int i, int j, int k) {arr[i] += k;if (j + 1 < arr.size()) arr[j + 1] -= k;}const vector<T>& query() const { return arr; }
};

差分数组类的打印:

ostream& operator<<(ostream& os, const DiffArray<int>& v) {return os << v.query();
}

测试:

int main() {//vector<int> v{4, 7, 5, 2, 9};cout << v;DiffArray<int> d1(v);cout << d1;d1.modify(2, 3, 1); // [2, 3] 增加 1cout << d1;cout << d1.get_raw();
}

结果:

4 7 5 2 9    // 原始数组
4 3 -2 -3 7  // 差分数组
4 3 -1 -3 6  // 修改后的差分数组
4 7 6 3 9    // 修改后的原数组

力扣题目

1109. 航班预订统计;

class DiffArray {vector<int> arr;public:DiffArray() : arr() {}DiffArray(vector<int> _arr) : arr(_arr) {}vector<int> get_raw() {vector<int> raw_arr(arr);for (int i{1}; i < arr.size(); ++i)raw_arr[i] = arr[i] + raw_arr[i - 1];return raw_arr;}void modify(int i, int j, int k) {arr[i] += k;if (j + 1 < arr.size()) arr[j + 1] -= k;}
};
class Solution {
public:vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {vector<int> ans(n);DiffArray da(ans);for (auto v : bookings) da.modify(v[0] - 1, v[1] - 1, v[2]);return da.get_raw();}
};

当然, 可以更快一些:

class Solution {
public:vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {vector<int> ans(n);for (auto v : bookings) {auto R{v[1]}, val{v[2]};ans[v[0] - 1] += val;if (R < n) ans[R] -= val;}for (int i{1}; i < n; ++i) ans[i] = ans[i] + ans[i - 1];return ans;}
};

最快的方法STL 以及开辟(n+1)长度的数组减少判断:

class Solution {
public:vector<int> corpFlightBookings(vector<vector<int>>& bookings, int n) {vector<int> ans(n + 1);for (vector<int>& i : bookings) {ans[i[0] - 1] += i[2];ans[i[1]] -= i[2];}partial_sum(ans.begin(), ans.end(), ans.begin());ans.pop_back();return ans;}
};