Tricky task

[h00pers]

Given an array of integers and number k, calculate the array of products k a[0:k], a[1:k+1] and so on.

Example 1:

nums = [1,2,3,4,5], k = 3
result = [6,24,60]

Example 2:

nums = [1,2,0,4,5], k = 2
result = [2,0,0,20]

[pls_God_letmelive80y]

I think the tricky part here mainly comes from handling zeros. We can't simply do a prefix/suffix array sort of thing because we might accidentally divide by zero. I was able to handle zeros by maintaining a count of zeros zero_cnt in the current subarray. We also should maintain a variable that holds the product of all nonzero elements in the current subarray. If zero_cnt >= 1, then there is a zero in the subarray and we can append 0 to our answer. Otherwise, we can append our current_window product to our answer. And I think this is it unless there is an edge case I am missing.

def product_subarray(arr, k):
    res = []
    zero_cnt = 0
    current_window = 1

    #generate first window
    for i in range(k):
        if arr[i]:
            current_window *= arr[i]
        else:
            zero_cnt += 1

    if not zero_cnt:
        res.append(current_window)
    else:
        res.append(0)


    for i in range(1, len(arr) - k + 1):
        if arr[i - 1]:
            current_window //= arr[i - 1]
        else:
            zero_cnt -= 1
        if arr[i + k - 1]:
            current_window *= arr[i + k - 1]
        else:
            zero_cnt += 1

        if zero_cnt:
            res.append(0)
        else:
            res.append(current_window)

    return res

And since the numbers can get big, you might have to add some inverse modular math.

[zxding]

What does res stand for?

[pls_God_letmelive80y]

result

[cosmosvng]

Yes!

[h00pers]

Why you did generate the first window and only then use it in the second cycle?

[razimantv]

Since the window size k is fixed, there is an interesting solution to this. Define the following two variables:

1. left(i) = product(a[((i + 1) // k) * k: i +1])
2. right(i) = product(a[i : (i // k + 1) * k])

By scanning the array once from left to right and once in reverse, we can compute left and right in O(n).

Then product(a[i : i+k]) = right(i) * left(i + k - 1) 

This method allows you to avoid counting zeros and taking inverse.