Google APAC 2016 Round E Problem D, Sums of Sums (aka APAC 2017 Practice Round Problem D)

D.py

#!/usr/bin/env python2
# Ref: https://www.quora.com/How-can-problem-D-Sums-of-sums-from-Round-E-of-the-Google-APAC-Test-2016-be-solved-for-the-large-dataset

import multiprocessing

get_int = lambda: int( raw_input() )
get_multi_int = lambda: map( int, raw_input().strip().split() )

def binary_search( x, f, lo, hi ):
    '''\
    find the smallest integer k in [ lo, hi ] that satisfies f(k) >= x
    '''
    while lo <= hi:
        mid = ( lo + hi ) / 2
        f_mid = f( mid )
        if f_mid < x:
            lo = mid+1
        else:
            hi = mid-1
    return lo

def preprocess( array ):
    '''\
    psum[i] = ar[1] + ar[2] + ... + ar[i]
    qsum[i] = N*ar[1] + (N-1)*ar[2] + ... + (N-i+1)*ar[i]
    '''
    global psum, qsum

    _sum = 0
    psum = [0]
    for elem in array:
        _sum += elem
        psum.append( _sum )

    _sum = 0
    qsum = [0]
    coeff = N
    for elem in array:
        _sum += coeff * elem
        qsum.append( _sum )
        coeff -= 1

def enumerate_subarray_sum_le( x ):
    R = 1
    for L in xrange( 1, N+1 ):
        while R <= N and psum[R]-psum[L-1] <= x: R += 1
        yield L, R-1

def count_subarray_sum_le( x ):
    total = 0
    for L, R in enumerate_subarray_sum_le( x ):
        total += R - L + 1
    return total

def sum_subarray_sum_le( x ):
    total = 0
    for L, R in enumerate_subarray_sum_le( x ):
        total += ( qsum[R] - qsum[L-1] ) \
              -  ( psum[R] - psum[L-1] ) * (N-R)
    return total

def subarray_sum_sorted_partial_sum( k ):
    ## step 1: find the value S of the k-th largest subarray sum
    S = binary_search( k, count_subarray_sum_le,
                       lo=0, hi=psum[-1] )
    real_k = count_subarray_sum_le( S )

    ## step 2: sum up all the subarrays with sum less than or equal to S
    return sum_subarray_sum_le( S ) - ( real_k - k ) * S

def query(( L, R )):
    ## 1 <= L <= R <= N*(N+1)/2
    return subarray_sum_sorted_partial_sum( R ) \
         - subarray_sum_sorted_partial_sum( L-1 )

def solve():
    global N

    ## input & preprocess
    N, Q = get_multi_int()
    array = get_multi_int()
    assert len( array ) == N
    preprocess( array )

    ## input & perform queries in parallel
    ## cpu_count()/2 assumes Hyper-Threading
    pool = multiprocessing.Pool( multiprocessing.cpu_count()/2 )
    tasks = ( get_multi_int() for _ in xrange( Q ))
    for result in pool.imap( query, tasks ):
        print result
    del pool

def main():
    T = get_int()
    for case_no in xrange( T ):
        print 'Case #%d:' % ( case_no+1 )
        solve()

if __name__ == '__main__':
    main()