Mistake in Python code

Question

There is an array of integers. There are also disjoint sets, A and B, each containing integers. You like all the integers in set A and dislike all the integers in set B. Your initial happiness is 0. For each integer in the array, if i in A, you add 1 to your happiness. If i in B, you add -1 to your happiness. Otherwise, your happiness does not change. Output your final happiness at the end.

Input Format

The first line contains integers n and m separated by a space. The second line contains n integers, the elements of the array. The third and fourth lines contain m integers, A and B respectively.

Output Format

Output a single integer, your total happiness.

Sample Input

3 2

1 5 3

3 1

5 7

Sample Output

1

Can someone please explain what is wrong with this solution? It passes some test, but fails on others.

input()
array = set(input().split())
set1 = set(input().split())
set2 = set(input().split())
res = len(set1 & array) - len(set2 & array)
print(res)

• francisco sollima: Can you provide a failing example so it'll be easier to spot the mistake?
• TemporalWolf: Probably because you're using sets: if I put in 1 3 3 5 you'd say that's 1 and it should be 2: the array cannot be reduced to a set.

Answer 1

The problem is that you're transforming your inputs to sets, which in turn removes the duplicates. If you have repeated values in your input, with the set you're only adding/substracting 1 to the resulting happiness. If that's correct, your code is fine. If not, then you should work with lists rather than sets.

The code could be something like this:

# The first part should stay the same, without the set() call on array
input()
array = input().split()
list1 = set(input().split())
list2 = set(input().split())
# Now we use some list comprehension to get the happiness result
res = sum([1 for elem in array if elem in list1]) - sum([1 for elem in array if elem in list2])

The first sum accumulates the positive points, and the second one the negatives. It works with multiple occurences, adding/substracting one point per each.

EDIT

A more clear approach, to understand the for loop

# The first part should stay the same, without the set() call on array
input()
array = input().split()
list1 = set(input().split())
list2 = set(input().split())
# We create a variable res which will store the resulting happiness, initially 0
res = 0
# Now we iterate through the elements in array and check wheter they should add or substract
for elem in array:
 # If the element is in list1, we add 1 to res
 if elem in list1:
 res += 1
 # If the element is in list2, we substract 1 from res
 elif elem in list2:
 res -= 1

• TemporalWolf: list1 and list2 can stay sets, and in is much faster this way. additionally, you can combine the list comprehension if you want to only iterate over array once: sum([1 if elem in list1 else -1 if elem in list2 else 0 for elem in array]), although I would suggest just breaking it out if in real code, as it's not obvious what you are doing.
• francisco sollima: Yours is a better one-liner, but you're right, I better just break it out.

Answer 2

I took the inputs for list A and B and as a general list. I wanted to obtain happiness in 1 line using list comprehension as below. After I merged print and "happiness =" , in one line. Apparently, this is the solution to make the code faster.

input()
my_array = input().split()
listA=list(input().split())
listB=list(input().split())
print (sum(1 for data in my_array if data in listA)+sum(-1 for data in my_array if data in listB))