Python - Matching two non-decreasing arrays

Question

I am working on the following question: Note: The code is to work under a time limit of 0.4 seconds, which requires a O(N+M) solution Your friend have given you a list of n positive integers: a0, a1, a2, ... an-1, in increasing order.

Now, your friend is going to ask you m questions, each of the form, "Here is a positive integer B. Is B a part of the list? (B is in non-decreasing order) "

Your task is to answer these questions. Output a single integer, which is the number of times you say "YES".

Input Format: The first line of input will contain the integer n.

The second line of input will contain n positive integers, which is the list a, in increasing order.

The next line of input will contain the integer m, which is the number of questions.

The next m lines of input will each contain an integer B in non-decreasing order , which is the positive integer that your friend is asking you about.

Output Format: Output should consist of a single integer, which is the number of times you say "YES".

Here is my code:

n = int(raw_input())
a = [int(x) for x in raw_input().split()]
m = int(raw_input())
b = []
for y in range(0, m):
    b.append(int(raw_input()))

answer = 0
i = 0
j = 0
z = 1

while (i < n) and (j < m):
    if a[i] == b[j]:
        answer = answer + 1
        while ((z + j) < m):
            if a[i] == b[j + z]:
                answer = answer + 1
                z = z + 1
            elif a[i] != b[j + z]:
                i = i + 1
                j = j + z + 1
    elif a[i] < b[j]:
        i = i + 1
    elif a[i] > b[j]:
        j = j + 1

print answer

As the code may suggest, I am attempting to use a method that compares items in array A and array B, and if they don't match, add one to the index of the array with the smaller item. However, the code is not outputting the correct answer, and sometimes will loop on forever asking me for unlimited number of inputs.

I am a beginner at programming and I understand this is a very fundamental algorithm. I apologise in any statements not made clear or technical enough.

Thank you.

Answer 1

I wouldn't use a while loop here. That's why its running for infinite time. Just use a for loop. Also there are some very nice features in the python language that allow you to do this more quickly than iterating over the entire list yourself.

In [6]: count = 0

In [7]: a = [1, 2, 3, 4, 5, 100]

In [8]: b = [1, 3, 1000, -10, 12]

In [9]: for val in b:
   ...:     if val in a:
   ...:         count += 1
   ...:

In [10]: print count
2

Julien Bernu has a more elegant idea:

a = set([1, 2, 3, 4, 5, 100])
b = set([1, 3, 1000, -10, 12])
print len(a.intersection(b))

Answer 2

You can first keep track of how many times each element appear in b in a dictionary:

counts = {x: b.count(x) for x in b}
print sum(counts[x] for x in set(a).intersection(b))

Answer 3

Thank you everyone for the helps. @Greg @JulienBernu I improve my code on my original logic and solved the problem. The solution is in O(N+M) and works within 0.25 seconds for all inputs.

n = int(raw_input())
a = [int(x) for x in raw_input().split()]
m = int(raw_input())
b = []
for y in range(0, m):
    b.append(int(raw_input()))

answer = 0
i = 0
j = 0

while (i < n) and (j < m):
    if a[i] == b[j]:
        answer = answer + 1
        j = j + 1
    elif a[i] < b[j]:
        i = i + 1
    elif a[i] > b[j]:
        j = j + 1

print answer

Or for answer to all values when b is not in non-decreasing order, check if the list b is sorted using all(b[i] <= b[i + 1] for i in xrange(m-1)) . If it is not sorted, use the method of binary search which I have previously used as a solution with O(Mlog(N)).