The TLE solution of the third problem of Kick Start 2019 Round B.

diverse_subarray.py

from collections import defaultdict, deque


def max_trinket(s, a):
    '''
    Attributes:
    s: The maximum number of trinkets allowed of a single type
    a: Array of the types of the trinkets
    '''
    if not s:
        return 0
    
    # Find the least available number of nodes of the segment tree
    mask = 1
    while mask < len(a):
        mask <<= 1
    # Allocate the segment tree
    tree = [None]*(mask<<1)

    # Array of events. It only contains 0, 1, -s
    diff = [0]*len(a)

    # Record the positions of each type of the trinkets
    pos = defaultdict(deque)
    for i in range(len(a)):
        cnt = len(pos[a[i]])
        pos[a[i]].append(i)
        if cnt > s:
            diff[i] = 0
        elif cnt == s:
            diff[i] = -s
        else:
            diff[i] = 1

    def update_node(i):
        '''
        Update the internal node of which index is `i`.
        '''
        lc = tree[i << 1]
        rc = tree[(i << 1) + 1]
        '''
        tree[i][0]: The sum of the elements covered by the node
        tree[i][1]: The maximum sum of any prefix of elements that are covered by the node
        '''
        tree[i] = (lc[0] + rc[0], max(lc[1], lc[0] + rc[1]))

    def build(l=0, r=len(diff), i=1):
        '''
        Build the segment tree 
        '''
        if r - l == 1:
            tree[i] = (diff[l], diff[l])
            return
        m = (l + r + 1) >> 1
        build(l, m, i << 1)
        build(m, r, (i << 1) + 1)
        update_node(i)
    
    def update(pos, val, l=0, r=len(diff), i=1):
        '''
        Update the element in the segment tree
        '''
        if r - l == 1:
            tree[i] = (val, val)
            return
        m = (l + r + 1) >> 1
        if pos < m:
            update(pos, val, l, m, i << 1)
        else:
            update(pos, val, m, r, (i << 1) + 1)
        update_node(i)

    build()
    
    max_t = 0
    for i in range(len(a)):
        '''
        For example:
        Assume the events of a single type listed as the following:
        index: i i+a i+b i+c i+d  (0 < a < b < c < d)
        events 1   1  -s   0   0

        After finishing updating the maximum sum of any prefix of elements
        that are started from the ith event, the events of the type should
        be changed to:
        index: i i+a i+b i+c i+d  (0 < a < b < c < d)
        events 0   1   1  -s   0
        '''
        max_t = max(max_t, tree[1][1])
        update(i, 0)
        typ = a[i]
        if len(pos[typ]) > s:
            update(pos[typ][s], 1)
        if len(pos[typ]) > s + 1:
            update(pos[typ][s + 1], -s)
        pos[typ].popleft()
    return max_t

test = int(input())
for case in range(1, test + 1):
    n, s = tuple(int(x) for x in input().split(' '))
    a = [int(x) for x in input().split(' ')]
    ans = max_trinket(s, a)
    print('Case #{}: {}'.format(case, ans))