O(n) time complexity solution
This is my first time posting a solution, so I would be grateful for any feedback. As far as I can tell, the solutions posted so far are O(n^2) or worse time co
This is my first time posting a solution, so I would be grateful for any feedback. As far as I can tell, the solutions posted so far are O(n^2) or worse time complexity, and as the solution I found was O(n) I thought it worth sharing.\n\n# Intuition\n<!-- Describe your first thoughts on how to solve this problem. -->\nWe shift our perspective of this problem by (instead of considering the sum of scores of the subarrays) considering how many subarrays to which a given element of nums will contribute (i.e. how many subarrays containing the given element exist such that when the subarray is sorted, there exists an element directly to the left of the given element that is more than one less than the given element). Given that the contribution always has value one, summing these counts over all elements of nums will give the solution.\n\n# Approach\n<!-- Describe your approach to solving the problem. -->\nWe assume without loss of generality that the subarray sorts are stable (i.e. does not change the order of elements of equal value). Let n be the length of nums.\n\nThe necessary and sufficient conditions for the element of nums at index i to contribute to the score for a given subarray spanning from indices j1 to j2 inclusive (where 0 <= j1 <= j2 < n) are:\n- j1 <= i <= j2\n- nums[j] != nums[i] and nums[j] != nums[i] - 1 for all j1 <= j < i\n- nums[j] != nums[i] - 1 for all i < j <= j2 (using the stability of the sort)\n- There exists j1 <= j <= j2 such that nums[j] < nums[i]\n\nSuppose there exists 0 <= k < i such that nums[k] = nums[i] or nums[k] = nums[i] - 1. Then for any subarray such that j1 <= k, the second of these conditions is violated. Thus, any subarray for which the element at i contributes must have j1 > k. It further follows that if k1 is the largest index such that 0 <= k1 < i and nums[k1] = nums[i] or nums[k1] = nums[i] - 1, or -1 if no such index exists, then in order for the element at index i to contribute, j1 > k1.\n\nSimilarly, if k2 is the smallest index such that i < k2 < n and nums[k2] = nums[i] -1, or n if no such index exists, then in order for the element at index i to contribute, j2 < k2.\n\nNow consider the converse. Suppose we have indices j1 and j2 such that k1 < j1 <= i <= j2 < k2 (where k1 and k2 are defined as previous). By the definitions of k1 and k2 respectively, there does not exist k1 < j < i such that nums[j] = nums[i] or nums[j] = nums[i] - 1 and there does not exist i < j < k2 such that nums[j] = nums[i] - 1. Therefore, given that k1 < j1 and k2 > j2 we can conclude that there does not exist j1 <= j < i such that nums[j] = nums[i] or nums[j] = nums[i] - 1 and there does not exist i < j <= j2 such that nums[j] = nums[i] - 1.\n\nConsequently, the first three conditions for the element of nums at index i to contribute to the score for a given subarray spanning from j1 to j2 inclusive (where 0 <= j1 <= j2 < n) if and only if k1 < j1 <= i <= j2 < k2.\n\nWe now need to find the conditions for which, of these possibilities, the final condition is satisfied. Let m1 be the largest index such that k1 < m1 < i and nums[m1] < nums[i], or k1 if no such index exists. Similarly let m2 be the smallest index such that i < m2 < k2, or k2 if no such index exists. If k1 < j1 <= m1 or m2 <= j2 < k2 then there exists j1 <= j <= j2 such that nums[j] < nums[i] (namely either j = m1 or j = m2 depending on the condition satisfied), while conversely, if m1 < j1 <= i <= j2 < m2 then, by the definition of m1 and m2 there does not exist m1 < j < m2 such that nums[j] < nums[i] and so there does not exist j1 <= j <= j2 such that nums[j] < nums[i].\n\nWe can therefore conclude that the element of nums at index i contributes to the score for a given subarray spanning from j1 to j2 inclusive (where 0 <= j1 <= j2 < n) if and only if k1 < j1 <= i <= j2 < k2 and j1 <= m1 or j2 >= m2, where k1, k2, m1 and m2 are as previously defined.\n\nWe can calculate the total number of such combinations of j1 and j2 for given i relatively easily given that the definitios of k1, k2, m1 and m2 only depend on i and not on j1 or j2. We observe that the total count for given i is simply:\n\n(number of subarrays for which k1 < j1 <= i and i <= j2 < k2) - (number of subarrays for which m1 < j1 <= i and i <= j2 < m2)\n\nSince in both these terms the possible ranges of j1 and j2 do not allow j1 > j2, the number of subarrays for both terms is the product of the number of values available to j1 multiplied by the number of values available to j2. As such, the number of subarrays for which k1 < j1 <= i and i <= j2 < k2 is just (i - k1) * (k2 - i) and similarly the number of subarrays for which m1 < j1 <= i and i <= j2 <= m2 is just (i - m1) * (m2 - i). Putting all this together and recalling that m1 >= k1 and m2 <= k2 (which implies that every j1 and j2 satisfying m1 < j1 <= i2 and i <= j2 < m2 also satisfies k1 < j1 <= i2 and i <= j2 < k2) gives the contribution for the element of nums at index i:\n\n    Total contribution for element of nums at index i = (i - k1) * (k2 - i) - (i - m1) * (m2 - i)\n\nWe can therefore find the overall solution by for each and every i identifying k1, k2, m1 and m2 and calculating the above value. The solution is then the sum of these values.\n\nThis can be accomplished in 2 passes. During an initial left to right pass, k1 is found for each i by maintaining a dictionary recording the index at which each value was last encountered (the dictionary latest in the code below), while m1 is found for each i by maintaining a strictly increasing monotonic stack (the list smaller_stk in the code). These values for k1 and m1 for each index i are stored in lists (k1 in k1_lst and the m1 m1_lst in the code). k2 and m2 can be similarly found for each i during a right to left pass, but on this occasion do not need to be stored given that we now have all the parameters needed to calculate the contribution to the final result for each element of nums as we encounter it. We therefore simply calculate that contribution using the formula derived above during this second pass and add it to the end result. This total gives the final answer.\n\nRegarding time complexity, as we are only iterating over the list twice and during each iteration adding to and reading no more than 2 values to a dictionary (O(1)), adding one element to a stack (O(1)) and removing elements from that stack (overall O(n) over all iterations), as well as other operations that are considered O(1), this algorithm has overall O(n) time complexity.\n\nRegarding memory complexity, we are maintaining a fixed number of dictionaries, lists and stacks that contain only numbers and are at most length n, so the space complexity is also O(n)\n\nThank you for taking the time to read this solution, I hope it was useful and understandable. I apologise if it was confusing or difficult to follow, I tried to be as clear as possible but that wasn\'t easy for this particular problem and I wanted to publish reasonably quickly. Again, as this is my first posted solution I would appreciate any feedback.\n\n# Complexity\n- Time complexity:\n<!-- Add your time complexity here, e.g. $$O(n)$$ -->\nO(n)\n\n- Space complexity:\n<!-- Add your space complexity here, e.g. $$O(n)$$ -->\nO(n)\n\n# Code\n```\nclass Solution:\n    def sumImbalanceNumbers(self, nums: List[int]) -> int:\n        # O(n) time complexity and O(n) space complexity\n        n = len(nums)\n        latest = {}\n        smaller_stk = []\n        m1_lst = [0] * n\n        k1_lst = [0] * n\n        for i, num in enumerate(nums):\n            k1_lst[i] = max(latest.get(num, -1), latest.get(num - 1, -1))\n            while smaller_stk and smaller_stk[-1][1] >= num:\n                smaller_stk.pop()\n            m1_lst[i] = max(smaller_stk[-1][0], k1_lst[i])\\\n                            if smaller_stk else k1_lst[i]\n            smaller_stk.append((i, num))\n            latest[num] = i\n        res = 0\n        soonest = {}\n        smaller_stk = []\n        for i in reversed(range(n)):\n            num = nums[i]\n            k2 = soonest.get(num - 1, n)\n            while smaller_stk and smaller_stk[-1][1] >= num:\n                smaller_stk.pop()\n            m2 = min(smaller_stk[-1][0], k2) if smaller_stk else k2\n            smaller_stk.append((i, num))\n            soonest[num] = i\n            res += (i - k1_lst[i]) * (k2 - i) -\\\n                    (i - m1_lst[i]) * (m2 - i)\n        return res\n            \n```
['Python3']