2281. Sum of Total Strength of Wizards Hard

1/**
2 * [2281] Sum of Total Strength of Wizards
3 *
4 * As the ruler of a kingdom, you have an army of wizards at your command.
5 * You are given a 0-indexed integer array strength, where strength[i] denotes the strength of the i^th wizard. For a contiguous group of wizards (i.e. the wizards' strengths form a subarray of strength), the total strength is defined as the product of the following two values:
6 *
7 * The strength of the weakest wizard in the group.
8 * The total of all the individual strengths of the wizards in the group.
9 *
10 * Return the sum of the total strengths of all contiguous groups of wizards. Since the answer may be very large, return it modulo 10^9 + 7.
11 * A subarray is a contiguous non-empty sequence of elements within an array.
12 *  
13 * Example 1:
14 *
15 * Input: strength = [1,3,1,2]
16 * Output: 44
17 * Explanation: The following are all the contiguous groups of wizards:
18 * - [1] from [<u>1</u>,3,1,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
19 * - [3] from [1,<u>3</u>,1,2] has a total strength of min([3]) * sum([3]) = 3 * 3 = 9
20 * - [1] from [1,3,<u>1</u>,2] has a total strength of min([1]) * sum([1]) = 1 * 1 = 1
21 * - [2] from [1,3,1,<u>2</u>] has a total strength of min([2]) * sum([2]) = 2 * 2 = 4
22 * - [1,3] from [<u>1,3</u>,1,2] has a total strength of min([1,3]) * sum([1,3]) = 1 * 4 = 4
23 * - [3,1] from [1,<u>3,1</u>,2] has a total strength of min([3,1]) * sum([3,1]) = 1 * 4 = 4
24 * - [1,2] from [1,3,<u>1,2</u>] has a total strength of min([1,2]) * sum([1,2]) = 1 * 3 = 3
25 * - [1,3,1] from [<u>1,3,1</u>,2] has a total strength of min([1,3,1]) * sum([1,3,1]) = 1 * 5 = 5
26 * - [3,1,2] from [1,<u>3,1,2</u>] has a total strength of min([3,1,2]) * sum([3,1,2]) = 1 * 6 = 6
27 * - [1,3,1,2] from [<u>1,3,1,2</u>] has a total strength of min([1,3,1,2]) * sum([1,3,1,2]) = 1 * 7 = 7
28 * The sum of all the total strengths is 1 + 9 + 1 + 4 + 4 + 4 + 3 + 5 + 6 + 7 = 44.
29 *
30 * Example 2:
31 *
32 * Input: strength = [5,4,6]
33 * Output: 213
34 * Explanation: The following are all the contiguous groups of wizards:
35 * - [5] from [<u>5</u>,4,6] has a total strength of min([5]) * sum([5]) = 5 * 5 = 25
36 * - [4] from [5,<u>4</u>,6] has a total strength of min([4]) * sum([4]) = 4 * 4 = 16
37 * - [6] from [5,4,<u>6</u>] has a total strength of min([6]) * sum([6]) = 6 * 6 = 36
38 * - [5,4] from [<u>5,4</u>,6] has a total strength of min([5,4]) * sum([5,4]) = 4 * 9 = 36
39 * - [4,6] from [5,<u>4,6</u>] has a total strength of min([4,6]) * sum([4,6]) = 4 * 10 = 40
40 * - [5,4,6] from [<u>5,4,6</u>] has a total strength of min([5,4,6]) * sum([5,4,6]) = 4 * 15 = 60
41 * The sum of all the total strengths is 25 + 16 + 36 + 36 + 40 + 60 = 213.
42 *
43 *  
44 * Constraints:
45 *
46 * 1 <= strength.length <= 10^5
47 * 1 <= strength[i] <= 10^9
48 *
49 */
50pub struct Solution {}
51
52// problem: https://leetcode.com/problems/sum-of-total-strength-of-wizards/
53// discuss: https://leetcode.com/problems/sum-of-total-strength-of-wizards/discuss/?currentPage=1&orderBy=most_votes&query=
54
55// submission codes start here
56
57impl Solution {
58    const MOD: i64 = 1_000_000_007;
59    pub fn total_strength(strength: Vec<i32>) -> i32 {
60        let l = strength.len();
61        let mut prefix = vec![0_i64; l + 1];
62        let mut prefix_sum = vec![0_i64; l + 2];
63
64        for i in 0..l {
65            prefix[i + 1] = (prefix[i] + strength[i] as i64) % Self::MOD;
66        }
67
68        for i in 0..l + 1 {
69            prefix_sum[i + 1] = (prefix_sum[i] + prefix[i]) % Self::MOD;
70        }
71
72        let mut left = vec![-1_i32; l];
73        let mut stack: Vec<usize> = vec![];
74        for i in 0..l {
75            while !stack.is_empty() && strength[stack[stack.len() - 1]] >= strength[i] {
76                stack.pop();
77            }
78            left[i] = if stack.is_empty() {
79                -1
80            } else {
81                stack[stack.len() - 1] as i32
82            };
83            stack.push(i);
84        }
85        let mut right = vec![l as i32; l];
86        stack.clear();
87        for i in (0..l).rev() {
88            while !stack.is_empty() && strength[stack[stack.len() - 1]] > strength[i] {
89                stack.pop();
90            }
91            right[i] = if stack.is_empty() {
92                l as i32
93            } else {
94                stack[stack.len() - 1] as i32
95            };
96            stack.push(i);
97        }
98
99        let mut result = 0;
100
101        for i in 0..l {
102            result += ((prefix_sum[right[i] as usize + 1] - prefix_sum[i + 1])
103                * (i as i64 - left[i] as i64)
104                % Self::MOD
105                + 2 * Self::MOD
106                - (prefix_sum[i + 1] - prefix_sum[(left[i] + 1) as usize])
107                    * (right[i] as i64 - i as i64)
108                    % Self::MOD)
109                % Self::MOD
110                * strength[i] as i64
111                % Self::MOD;
112            result = result % Self::MOD
113        }
114        result as i32
115    }
116}
117
118// submission codes end
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_2281() {
126        assert_eq!(Solution::total_strength(vec![1, 3, 1, 2]), 44);
127        assert_eq!(Solution::total_strength(vec![5, 4, 6]), 213);
128        assert_eq!(
129            Solution::total_strength(vec![
130                747, 812, 112, 1230, 1426, 1477, 1388, 976, 849, 1431, 1885, 1845, 1070, 1980, 280,
131                1075, 232, 1330, 1868, 1696, 1361, 1822, 524, 1899, 1904, 538, 731, 985, 279, 1608,
132                1558, 930, 1232, 1497, 875, 1850, 1173, 805, 1720, 33, 233, 330, 1429, 1688, 281,
133                362, 1963, 927, 1688, 256, 1594, 1823, 743, 553, 1633, 1898, 1101, 1278, 717, 522,
134                1926, 1451, 119, 1283, 1016, 194, 780, 1436, 1233, 710, 1608, 523, 874, 1779, 1822,
135                134, 1984
136            ]),
137            471441678
138        );
139    }
140}
141