Problem statement

https://binarysearch.com/problems/Non-Overlapping-Pairs-of-Sublists/

Solution

The idea is to first create groups of elements which are >= k. Say we have lengths [A1, ... Ak]. Then we can have two options:

  1. Choose sublists from different groups, it will be $C_{x+1}^2 \cdot C_{y+1}^2$ and we need to sum it for all pairs.
  2. Choose sublists from the same group, then we have $C_{x+1}^4$ options for choose ends of sublists.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, arr, k):
        A = []
        for x, y in groupby(arr, key = lambda x: x >= k):
            if x: A += [len(list(y))]
        
        F = lambda x: (x-1)*x*(x+1)*(x+2)//24
        G = lambda x: (x+1)*x//2

        part1 = (sum(G(x) for x in A)**2 - sum(G(x)**2 for x in A))//2
        part2 = sum(F(x) for x in A)
        return (part1 + part2) % (10**9 + 7)