Finding an element in an array that occurs a given number of times

Question

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

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 at least equal to f.
• Type 1: For this query, Fredo has to answer the first number in the array such that frequency 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 \le N \le 10^6\$
• \$1 \le A[i] \le 10^{18}\$
• \$1 \le Q \le 10^6\$
• \$0 \le \text{type} \le 1\$
• \$1 \le f \le 10^{18}\$

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

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])
    temp1 = list(key)
    temp2 = list(value)
    for w in range(t):
        try:
            index = temp1.index(a[w])
        except:
            continue
        if query == 0:
            if temp2[index] >= f:
                print a[w]
                break
            else:
                del temp1[index]
                del temp2[index]
        else:
            if temp2[index] == f:
                print a[w]
                break
            else:
                del temp1[index]
                del temp2[index]
        if len(temp1) == 0:
            print 0
            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

Style

Codding

I suggest you to make use of white space so your code is more readable; as it stand, it seems like a bunch of thoughts thrown away without even trying to organize them. Read PEP 8 too, which will make you code look a bit more like Python to people that are used to read some Python code.

It is also good practice to organize your code into functions, so you can extract out logical steps, name them and, as such, make it clearer to everyone what it is you are trying to accomplish. Read about PEP 257 to make it even more clear.

At the very least, if you can't find out good functions, you may want to put all your top-level code into a main function and have, at the bottom of your script:

if __name__ == '__main__':
    main()

You can read a bit more about this idiom in this SO post.

Reading input

In challenges like that, where the input is highly regularly formatted, I have a tendency to favor raw_input rather than manipulating stdin directly.

It lets you more easily parse each line independently; for instance:

def parse_data():
    raw_input()  # Skip the first line containing number of elements
    items = raw_input()
    return items.split()


def run_queries(data):
    count = int(raw_input())
    for _ in range(count):
        type_, element = raw_input().split()
        result = compute_query(int(type_), element, data)
        print(result)


def compute_query(type_, item, data):
    # TODO


if __name__ == '__main__':
    data = parse_data()
    run_queries(data)

You will note that I don't try to convert the elements to integers as there is no added value doing so. Equality of two strings is as valid than equality of two integers to know that two elements are the same.

Algorithm

You know about Counter which is a good way to count elements from an iterable. You can use that at your advantage to reverse the keys and values so you know, for each "number of repetitions", which items corresponds.

from collections import Counter


def count_elements(array):
    counts = Counter(array)

    indexor = {}
    for element, count in counts.iteritems():
        indexor.setdefault(count, []).append(element)

    return indexor

There are several improvements to this snippet. First of, indexor will only let you answer queries of type 1 where you need the exact amount of elements. You can easily create a second dictionary to answer type 0 queries where you decrease count and append to each entry in the range \$[0, count]\$. The second improvement is to use a defaultdict(list) from collections to avoir the call to setdefault.

But the main thing is: this snippet can't answer the question as Counter is a dictionnary and thus doesn't take the order of insertions into account. Luckily, collections also provides OrderedDict which do. So your best friend for this challenge will most likely be:

from collections import Counter, OrderedDict

class OrderedCounter(Counter, OrderedDict):
    pass

Replace the use of Counter by OrderedCounter and voilà, you're done. Each query can be simplified to getting the right dictionnary based on the query type and then getting the first element of the list accessed by the element count.

Putting all that together

from collections import Counter, OrderedDict, defaultdict


class OrderedCounter(Counter, OrderedDict):
    pass


def count_elements(array):
    counts = OrderedCounter(array)

    strict_indexor = defaultdict(list)
    loose_indexor = defaultdict(list)

    for element, count in counts.iteritems():
        strict_indexor[count].append(element)
        for c in range(count, 0, -1):
            loose_indexor[c].append(element)

    return loose_indexor, strict_indexor


def parse_data():
    raw_input()  # Skip the first line containing number of elements
    items = raw_input()
    return items.split()


def run_queries(items_counts):
    queries_count = int(raw_input())
    for _ in range(queries_count):
        type_, number_of_elements = raw_input().split()
        result = compute_query(int(type_), number_of_elements, items_counts)
        print(result)


def compute_query(type_, count, items_counts):
    indexor = items_counts[type_]
    return indexor[count][0]  # Assume there is no query having no answer


if __name__ == '__main__':
    items = parse_data()
    counts_by_query_type = count_elements(items)
    run_queries(counts_by_query_type)

Answer 2

Your code is hard to understand, very hard. What is tokenizedInput[k+2*j+1]? What is tokenizedInput[k+2*j+2]? It's just outright confusing.

Instead I'd highly recommend that you change how you get the input. You know that data is held on different lines, and so you can split on each line. After this I'd split on any white space, and change the items to numbers. This simply becomes:

lines = [map(int, line.split()) for line in sys.stdin.read().split('\n')]

After this, you can extract the data, you know the first and third line will be an array on one. You know that the array is on the second line, and you know the rest are quires. This allows you to use:

(array_size,), array, (number_queries,) = lines[:3]
array = array[:array_size]
queries = lines[3:number_queries+3]

---

As for your code. Why on earth would you change \$O(1)\$ lookup to \$O(n)\$?! temp1.index(a[w]) is horrid. Instead actually use the benefits of a dictionary. I'd also change your variable names to be better than temp1. And you should never have bare excepts, as it can mask bugs. This could result in:

for query, f in queries:
    c = collection.copy()
    for w, item in enumerate(array):
        try:
            value = c[item]
        except KeyError:
            continue
        if query == 0:
            if value >= f:
                print item
                break
            else:
                del c[item]
        else:
            if value == f:
                print item
                break
            else:
                del c[item]

        if len(c) == 0:
            print 0
            break
        if w == array_size-1:
            print 0

However this still requires you to copy the dictionary collection. Instead change this code to be two generator comprehensions, where they work differently on the two different queries. And then use next to get the first occurrence. This should be a simple transform, to:

for query, freq in queries:
    if query == 0:
        items = (item for item in array if collection.get(item, -1) >= freq)
    else:
        items = (item for item in array if collection.get(item, -1) == freq)
    print next(items, 0)

You can then remove the duplicate code, by using the operator library.

Finally to get another speed up you should make array unique, so that if I enter the array 1 1 1 1 1 1 2 it becomes 1 2. This is as you don't need the other 1s as you store how many 1s there are in collection. This should get you code that looks like:

from collections import Counter
import sys
import operator

lines = [map(int, line.split()) for line in sys.stdin.read().split('\n')]
(array_size,), array, (number_queries,) = lines[:3]
array = array[:array_size]
queries = lines[3:number_queries+3]

collection = Counter(array)

new_array = []
already_seen = set()
for item in array:
    if item in already_seen:
        continue
    new_array.append(item)
    already_seen.add(item)
array = new_array

for query, f in queries:
    fn = operator.ge if query == 0 else operator.eq
    items = (item for item in array if fn(collection.get(item, -1), f))
    print next(items, 0)

This is \$O(k)\$ where yours was \$O(nk)\$. Where \$n\$ is the size of the array, and \$k\$ is the size of the unique array.