Python program to find happiness by checking participation of elements into sets

Suppose we have an array nums with n different integers. We also have two disjoint sets A and B. We have one happiness parameter which is set to 0 initially. We go through each integer i in nums. If i is in A then add happiness by 1 and if i is in B decrease it by 1. We have to finally find the final happiness value.

So, if the input is like nums = [1,2,5,8,6,3], A = {5,8,9,7,3}, B = {2,4,12,15}, then the output will be 2 because 5, 8, 3 are in A so happiness is 3 now, but 2 is in B so decrease it by 1 then happiness is 2.

Algorithm

To solve this, we will follow these steps −

• happiness := 0
• for each i in nums, do
  • if i is in A, then
    • happiness := happiness + 1
  • otherwise when i is in B, then
    • happiness := happiness - 1
• return happiness

Example

Let us see the following implementation to get better understanding −

def solve(nums, A, B):
    happiness = 0
    for i in nums:
        if i in A:
            happiness += 1
        elif i in B:
            happiness -= 1
    return happiness

nums = [1, 2, 5, 8, 6, 3]
A = {5, 8, 9, 7, 3}
B = {2, 4, 12, 15}
print(solve(nums, A, B))

The output of the above code is −

2

How It Works

The algorithm iterates through each element in the nums array. For element 1: not in A or B (happiness remains 0). For element 2: found in B (happiness decreases to -1). For element 5: found in A (happiness increases to 0). For element 8: found in A (happiness increases to 1). For element 6: not in A or B (happiness remains 1). For element 3: found in A (happiness increases to 2).

Alternative Using Set Operations

We can also solve this using set intersection operations −

def solve_optimized(nums, A, B):
    nums_set = set(nums)
    positive_count = len(nums_set & A)
    negative_count = len(nums_set & B)
    return positive_count - negative_count

nums = [1, 2, 5, 8, 6, 3]
A = {5, 8, 9, 7, 3}
B = {2, 4, 12, 15}
print(solve_optimized(nums, A, B))

2

Conclusion

The happiness calculation problem demonstrates efficient set membership operations in Python. The first approach iterates through the array, while the optimized version uses set intersections to count matching elements directly.