The following code is my solution to an old problem I found today.
Input Format
Line 1: Two integers, N (1 ≤ N ≤ 106) and window size K (1 ≤ K ≤ N)
Line 2: N positive integers of upvote counts, each integer less than or equal to 109Output Format
Line 1..: N - K + 1 integers, one integer for each window's result on each line
Sample Input
5 3 1 2 3 1 1
Sample Output
3 0 -2
Explanation
For the first window of [1, 2, 3], there are 3 non-decreasing subranges and 0 non-increasing, so the answer is 3. For the second window of [2, 3, 1], there is 1 non-decreasing subrange and 1 non-increasing, so the answer is 0. For the third window of [3, 1, 1], there is 1 non-decreasing subrange and 3 non-increasing, so the answer is -2.
I could solve the problem's logic and the initial test cases passed. But for the longer test cases, it gives timeout.
Are there any ways I can optimize parts of the following code?
def window_result(arr, size):
num_inc = 0
count = 0
for elem in arr:
if elem == 'i' or elem == 'e':
count += 1
else:
num_inc += (count*(count+1))/2
count = 0
num_inc += (count*(count+1))/2
num_dec = 0
count = 0
for elem in arr:
if elem == 'd' or elem == 'e':
count += 1
else:
num_dec += (count*(count+1))/2
count = 0
num_dec += (count*(count+1))/2
return num_inc-num_dec
conf = raw_input().strip().split()
num_days, window_size = int(conf[0]), int(conf[1])
data = [int(i) for i in raw_input().strip().split()]
begin = 0
end = window_size - 1
while end < num_days:
diff_array = []
for i in xrange(begin + 1, end + 1):
if data[i] > data[i - 1]:
diff_array.append('i')
elif data[i] == data[i - 1]:
diff_array.append('e')
else:
diff_array.append('d')
print window_result(diff_array, window_size-1)
begin += 1
end += 1