字節跳動 Magic

【題目描述】 給一個包含 n個整數元素的集合 個整數元素的集合 a，一個包含 m個整數元素的集合 b。
定義 magic操作爲，從一個集合中取出元素放到另裏且過後每的平均值都大於操作前。
注意以下兩點：
①不可以把一個集合的元素取空，這樣就沒有平均值了
②值爲 x的元素從集合 b取出放入集合a，但集合 a中已經有值爲 x的元素，則 a的平均值不變（因爲集 合元素不會重複）， b的平均值可能會改變（因爲 x被取出了）
問最多可以進行少次 magic操作
輸入：
3 5
1 2 5
2 3 4 5 6
輸出：
2

 

題目整體不難，記錄一下幾個細節點：

1.題目涉及轉移元素，因爲轉移後被轉的平均數要變大，而轉入相同的數不會替換原來的元素，也不會改變原本的平均值，所以這裏存放數據的時候都用的是集合，即set. set裏的元素就像字典的key值，不存在重複。集合用set()構建。因爲涉及後續比較，格式要轉爲float。

2.用到了一個python中的無窮大作爲比較。寫法：float('inf')。

自己的代碼如下：

class Magic:
    def findMin(self, biggerset, smallerset, meana, meanb):
        min_val = float('inf')
        if not biggerset:
            return -1

        for element in biggerset:
            if min_val > element and element < meana and element > meanb and element not in smallerset:
                min_val = element

        if min_val == float('inf'):
            return -1
        return min_val

    def swap(self, a, b):
        return [b, a]

    def mean(self, seta):
        return sum(seta)/len(seta)

if __name__ == '__main__':
    [n, m] = list(map(int, input().split(' ')))
    seta = set(map(float, input().split(' ')))  # 用set確定其不能重複
    setb = set(map(float, input().split(' ')))

    m = Magic()
    meana = m.mean(seta)
    meanb = m.mean(setb)
    if meana < meanb:  # 確保是seta的平均比setb大
        temp = set()
        temp = seta
        seta = setb
        setb = temp
        [meana, meanb] = m.swap(meana, meanb)

    step = 0
    min_val = m.findMin(seta, setb, meana, meanb)
    while(min_val != -1):
        seta.remove(min_val)
        meana = m.mean(seta)
        setb.add(min_val)
        meanb = m.mean(setb)
        step += 1
        min_val = m.findMin(seta, setb, meana, meanb)

    print(step)