pythonset

Mistake in Python code


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)

Solution

  • 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