How to improve the run time of a python code

Question

I am trying to solve the below problem from an online coding site.

QUESTION: Fredo is pretty good at dealing large numbers. So, once his friend Zeus gave him an array of N numbers , followed by Q queries which he has to answer. In each query , he defines the type of the query and the number f for which Fredo has to answer. Each query is of the following two types: Type 0: For this query, Fredo has to answer the first number in the array (starting from index 0) such that its frequency is atleast equal to f. Type 1: For this query, Fredo has to answer the first number in the array such that frequecy is exactly equal to f. Now, Fredo answers all his queries but now Zeus imagines how he should verify them . So, he asks you to write a code for the same. Note: If there is no number which is the answer to the query, output 0. Use fast I/O.

Input : The first line of the input contains N , the size of the array The next line contains N space separated integers. The next line contains Q, denoting the number of queries. Then follow Q lines, each line having two integers type and f, denoting the type of query and the frequency for which you have to answer the query.

Output: You have to print the answer for each query in a separate line.

Input Constraints:

1≤N≤106

1≤A[i]≤1018

1≤Q≤106

0≤type≤1

1≤f≤1018

SAMPLE INPUT

6

1 2 2 1 2 3

5

0 1

0 2

1 2

1 3

0 3

SAMPLE OUTPUT

1

1

1

2

2

Solution: Here is the solution that I have tried

from collections import Counter
import sys
tokenizedInput = sys.stdin.read().split()
t = int(tokenizedInput[0])
a = []
for i in range(t):
    s = int(tokenizedInput[i+1])
    a.append(s)
collection = Counter(a)
key = collection.keys()
value = collection.values()
q = int(tokenizedInput[i+2])
k = i+2
for j in range(q):
    query = int(tokenizedInput[k+2*j+1])
    f = int(tokenizedInput[k+2*j+2])
    for w in range(t):
        index = key.index(a[w])
        if query == 0:
            if value[index] >= f:
                print a[w]
                break
        else:
            if value[index] == f:
                print a[w]
                break
        if w == t-1:
            print 0

This code runs properly and gives the correct output for smaller test cases but crosses the time limit on larger test cases. Can someone please suggest what improvements can be made to this code to improve the speed.

Answer 1

Some suggestions: get the int() conversion of tokenizedInput done and out of the way rather than repeatedly calling int() over and over in the loop; a dictionary like collection is very efficient, don't second guess it by extracting the keys and values, use it as it was intended; precalculate anything you can before the loop; simplify, simplify, simplify.

I've reworked your code along the suggestions above, plus other tweaks, see if it makes sense to you and performs the exercise within the time limit:

import sys
from collections import Counter

tokenizedInput = map(int, sys.stdin.read().split())

N = tokenizedInput[0]

array = tokenizedInput[1:N + 1]

collection = Counter(array)

Q = tokenizedInput[N + 1]

k = N + 2

for j in range(Q):
    offset = 2 * j + k
    query, frequency = tokenizedInput[offset:offset + 2]

    for w in range(N):
        value = collection[array[w]]

        if query == 0:
            if value >= frequency:
                print array[w]
                break
        else:
            if value == frequency:
                print array[w]
                break

    else:  # aka "no break"
        print 0

Most folks would have handled this with separate read statements to input the various scalar and array values. You chose to do it in one read at the beginning, which is fine, but you must take care in your design that that initial choice doesn't get in your way later on.