Update sum-of-total-strength-of-wizards.py

Author: kamyu104

Python/sum-of-total-strength-of-wizards.py

@@ -20,9 +20,9 @@ def totalStrength(self, strength):
         stk, result = [-1], 0
         for i in xrange(len(strength)+1):
             while stk[-1] != -1 and (i == len(strength) or strength[stk[-1]] >= strength[i]):
-                # strength[j] >= strength[y] for j in [x, y]
-                # strength[j] < strength[y] for j in [y+1, z]
                 x, y, z = stk[-2]+1, stk.pop(), i-1
+                # assert(all(strength[j] >= strength[y] for j in xrange(x, y+1)))
+                # assert(all(strength[j] > strength[y] for j in xrange(y+1, z+1)))
                 result = (result+(strength[y]*((z-y+1)*((prefix2[y+1]-prefix2[x])-x*(prefix[y+1]-prefix[x]))+
                                                (y-x+1)*((suffix2[y+1]-suffix2[z+1])-(len(strength)-(z+1))*(suffix[y+1]-suffix[z+1])))))%MOD
             stk.append(i)