c++ 算法【前缀和】【单调栈】LeetCode2281:巫师的总力量和

硬件工程师的前景综合 2024-01-02 18 0

作者推荐

map|动态规划|单调栈|LeetCode975:奇偶跳

本文涉及的基础知识点

单调栈分类、封装和总结 C++算法：前缀和、前缀乘积、前缀异或的原理、源码及测试用例包括课程视频

题目

作为国王的统治者，你有一支巫师军队听你指挥。给你一个下标从 0 开始的整数数组 strength ，其中 strength[i] 表示第 i 位巫师的力量值。对于连续的一组巫师（也就是这些巫师的力量值是 strength 的子数组），总力量定义为以下两个值的乘积：巫师中最弱的能力值。组中所有巫师的个人力量值之和。请你返回所有巫师组的总力量之和。由于答案可能很大，请将答案对 109 + 7 取余后返回。子数组是一个数组里非空连续子序列。示例 1：输入：strength = [1,3,1,2] 输出：44 解释：以下是所有连续巫师组：

[1,3,1,2] 中 [1] ，总力量值为 min([1]) * sum([1]) = 1 * 1 = 1[1,3,1,2] 中 [3] ，总力量值为 min([3]) * sum([3]) = 3 * 3 = 9[1,3,1,2] 中 [1] ，总力量值为 min([1]) * sum([1]) = 1 * 1 = 1[1,3,1,2] 中 [2] ，总力量值为 min([2]) * sum([2]) = 2 * 2 = 4[1,3,1,2] 中 [1,3] ，总力量值为 min([1,3]) * sum([1,3]) = 1 * 4 = 4[1,3,1,2] 中 [3,1] ，总力量值为 min([3,1]) * sum([3,1]) = 1 * 4 = 4[1,3,1,2] 中 [1,2] ，总力量值为 min([1,2]) * sum([1,2]) = 1 * 3 = 3[1,3,1,2] 中 [1,3,1] ，总力量值为 min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5[1,3,1,2] 中 [3,1,2] ，总力量值为 min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6[1,3,1,2] 中 [1,3,1,2] ，总力量值为 min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7 所有力量值之和为 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44 。示例 2：

输入：strength = [5,4,6] 输出：213 解释：以下是所有连续巫师组：

[5,4,6] 中 [5] ，总力量值为 min([5]) * sum([5]) = 5 * 5 = 25[5,4,6] 中 [4] ，总力量值为 min([4]) * sum([4]) = 4 * 4 = 16[5,4,6] 中 [6] ，总力量值为 min([6]) * sum([6]) = 6 * 6 = 36[5,4,6] 中 [5,4] ，总力量值为 min([5,4]) * sum([5,4]) = 4 * 9 = 36[5,4,6] 中 [4,6] ，总力量值为 min([4,6]) * sum([4,6]) = 4 * 10 = 40[5,4,6] 中 [5,4,6] ，总力量值为 min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60 所有力量值之和为 25 + 16 + 36 + 36 + 40 + 60 = 213 。参数范围： 1 <= strength.length <= 105 1 <= strength[i] <= 109

单调栈+两个前缀和

时间复杂度: O(n)。

单调栈

枚举各数组的最小值。left 是从右向左第一个小于strength[i]的下标是left，如果不存在left是-1;right是从左向右第一个小于等于strength[i]的下标，如果不存在right为m_c。如果一个子数组strength[li,ri]的最小值是strength[i]，且strength[i]左边没有和它相等的值。则： li的取值范围是：(left,i] ri的取值范围是：[i,right) 我们可以这样计算这些子数组和的和。累加strength[i]它出现的次数。无论li和ri取值什么，i都在子数组中，所以i出现：(i-left)(right-i)次。 li为left+1时,ri无论取值什么，都包括strength[left+1] li为left+1和left+2时，ri无论取值什么，都包括strength[left+2] …

(left,i)中的元素分别出现:right-i次、(right-i)*2、(right-i)*3…

(i,right)中的元素，分别出现：…、(i-left)*2、(i-left)次

前缀和

vPreSum[i] 记录：strength[j]的和，j取值范围[0,i)。 vPreSum2[i]记录：strength[j]*(j+1)的和，j取值范围[0,i)。

代码

核心代码

template

class C1097Int

{

public:

C1097Int(long long llData = 0) :m_iData(llData% MOD)

{

}

C1097Int operator+(const C1097Int& o)const

{

return C1097Int(((long long)m_iData + o.m_iData) % MOD);

}

C1097Int& operator+=(const C1097Int& o)

{

m_iData = ((long long)m_iData + o.m_iData) % MOD;

return *this;

}

C1097Int& operator-=(const C1097Int& o)

{

m_iData = (m_iData + MOD - o.m_iData) % MOD;

return *this;

}

C1097Int operator-(const C1097Int& o)

{

return C1097Int((m_iData + MOD - o.m_iData) % MOD);

}

C1097Int operator*(const C1097Int& o)const

{

return((long long)m_iData * o.m_iData) % MOD;

}

C1097Int& operator*=(const C1097Int& o)

{

m_iData = ((long long)m_iData * o.m_iData) % MOD;

return *this;

}

bool operator<(const C1097Int& o)const

{

return m_iData < o.m_iData;

}

C1097Int pow(long long n)const

{

C1097Int iRet = 1, iCur = *this;

while (n)

{

if (n & 1)

{

iRet *= iCur;

}

iCur *= iCur;

n >>= 1;

}

return iRet;

}

C1097Int PowNegative1()const

{

return pow(MOD - 2);

}

int ToInt()const

{

return m_iData;

}

private:

int m_iData = 0;;

};

class CRangIndex

{

public:

template

CRangIndex(const vector& nums, _Pr CurIndexCmpStackTopIndex)

{

m_c = nums.size();

m_vLeft.assign(m_c, -1);

m_vRight.assign(m_c, m_c);

stack sta;

for (int i = 0; i < m_c; i++)

{

while (sta.size() && (CurIndexCmpStackTopIndex(i, sta.top())))

{

m_vRight[sta.top()] = i;

sta.pop();

}

if (sta.size())

{

m_vLeft[i] = sta.top();

}

sta.emplace(i);

}

int m_c;

vector m_vLeft, m_vRight;//vLeft[i] 从右向左第一个小于nums[i] ;vRight[i] 是第一个小于等于nums[i]。

};

template

vector CreatePreSum(const vector& nums)

{

vector preSum;

preSum.push_back(0);

for (int i = 0; i < nums.size(); i++)

{

preSum.push_back (preSum[i]+ nums[i]);

}

return preSum;

}

class Solution {

public:

int totalStrength(vector& strength) {

CRangIndex ri(strength, [&](int i1, int i2) {return strength[i1] <= strength[i2]; });

auto vPreSum = CreatePreSum>(strength);

vector> vPreSum2 = { 0 };

for (int i = 0 ; i < ri.m_c ; i++ )

{

const auto& n = strength[i];

vPreSum2.emplace_back(vPreSum2.back() + C1097Int<>(n)*(i+1));

}

C1097Int<> biRet = 0;

for (int i = 0; i < ri.m_c; i++)

{

const int left = ri.m_vLeft[i];

const int right = ri.m_vRight[i];

const int leftLen = i - left;//左边界的可选范围(left,i]

const int rightLen = right - i;//右边界的可选范围[i,right)

//计算(left,i)的巫师和

C1097Int<> leftTotal = vPreSum2[i] - vPreSum2[left + 1] - (vPreSum[i] - vPreSum[left + 1]) * (left+1);

//计算(i,right)的巫师和

C1097Int<> rightTotal = (vPreSum[right] - vPreSum[i + 1])*(right+1) - ( vPreSum2[right] - vPreSum2[i + 1] );

C1097Int<> iTotal = C1097Int<>(strength[i]) * leftLen * rightLen;

//计算以小标i为最弱力量的子数组

C1097Int<> cur = leftTotal * rightLen + rightTotal * leftLen + iTotal;

cur *= strength[i];

biRet += cur;

}

return biRet.ToInt();

}

};

测试用例

template

void Assert(const T& t1, const T& t2)

{

assert(t1 == t2);

}

template

void Assert(const vector& v1, const vector& v2)

{

if (v1.size() != v2.size())

{

assert(false);

return;

}

for (int i = 0; i < v1.size(); i++)

{

Assert(v1[i], v2[i]);

}

int main()

{

vector strength;

{

Solution slu;

strength = { 2,3,4 };