題目
三數(shù)之和
給定一個包含 n 個整數(shù)的數(shù)組 nums，判斷 nums 中是否存在三個元素 a，b，c ，使得 a + b + c = 0 ？找出所有滿足條件且不重復(fù)的三元組。
注意：答案中不可以包含重復(fù)的三元組。
例如, 給定數(shù)組 nums = [-1, 0, 1, 2, -1, -4]，
滿足要求的三元組集合為：
[
[-1, 0, 1],
[-1, -1, 2]
]
思路
首先想到的就是可以進行3個for循環(huán)，暴力直接解決
代碼如下:
func threeSum(_ nums: [Int]) -> [[Int]] {
var returnArray: [[Int]] = []
for num1 in nums {
for num2 in nums {
for num3 in nums {
if num1 + num2 + num3 == 0 {
returnArray.append([num1, num2, num3])
}
}
}
}
return returnArray
}
運行結(jié)果如下：
[[-1,-1,2],[-1,0,1],[-1,1,0],[-1,2,-1],[-1,2,-1],[-1,-1,2],[0,-1,1],[0,0,0],[0,1,-1],[0,1,-1],[0,-1,1],[1,-1,0],[1,0,-1],[1,0,-1],[1,-1,0],[2,-1,-1],[2,-1,-1],[2,2,-4],[2,-1,-1],[2,-1,-1],[2,-4,2],[-1,-1,2],[-1,0,1],[-1,1,0],[-1,2,-1],[-1,2,-1],[-1,-1,2],[-4,2,2]]
運行結(jié)果明顯與預(yù)期結(jié)果差別太大，一看就知道是條件判斷沒做好，導(dǎo)致數(shù)字重復(fù)使用。
怎么才能防止重復(fù)使用數(shù)字呢？于是想到可以先對數(shù)組進行排序，然后每次遍歷都取當(dāng)前 index 后面的 index 就可以防止重復(fù)遍歷一個元素。
但是這個方式還是會出現(xiàn)重復(fù)的數(shù)組，如在題目中給的范例 [-1, 0, 1]這個結(jié)果就會出現(xiàn)2次，因為排序后-1有2個，都可以和后面的0和1相加為0，所以還需要加上去重。我這里就是使用了 Set 來去重。插入時是插入 Set 自動去重，最后返回的時候轉(zhuǎn)換為 Array 返回。
修改后的代碼如下:
func threeSum(_ nums: [Int]) -> [[Int]] {
var numSet = Set<[Int]>()
let nums = nums.sorted()
for (index1, num1) in nums.enumerated() {
for (index2, num2) in nums[(index1 + 1)...].enumerated() {
for num3 in nums[(index1 + index2 + 2)...] {
if num1 + num2 + num3 == 0 {
numSet.insert([num1, num2, num3])
}
}
}
}
return Array(numSet)
}
上面的代碼雖然通過了執(zhí)行，但是提交以后的結(jié)果是失敗的，因為超出了時間限制，看來是暴力法太粗暴了，時間復(fù)雜度達到了 O(n3) ，所以還需要繼續(xù)優(yōu)化代碼。
于是我想到了之前 TwoSum 時，通過字典的 key 來保存相加需要的數(shù)字來判斷的辦法。這個優(yōu)化后的代碼應(yīng)該是可以達到 O(n2) 的時間復(fù)雜度的。于是就開始寫代碼，加一個字典，先存好所有數(shù)字作為 key，然后2個循環(huán)獲取數(shù)字。
優(yōu)化后的代碼如下：
func threeSum(_ nums: [Int]) -> [[Int]] {
var numSet = Set<[Int]>()
var numDictionary = [Int: Int]()
let nums = nums.sorted()
for (index, num) in nums.enumerated() {
numDictionary[num] = index
}
for (index1, num1) in nums.enumerated() {
for (index2, num2) in nums[(index1 + 1)...].enumerated() {
let num = 0-num1-num2
if let index = numDictionary[num] {
if index > index1 + index2 + 1 {
numSet.insert([num1, num2, num])
}
}
}
}
return Array(numSet)
}
結(jié)果還是沒有通過提交，還是超過了時間限制?？戳讼略斍?，是在[0,0 …… 0,0](幾百個0)這個用例中沒有執(zhí)行完成的。與是對這個情況進行針對性的優(yōu)化，判斷如果已經(jīng)有重復(fù)的數(shù)字處理過就不再進行處理。
最后代碼如下：
func threeSum(_ nums: [Int]) -> [[Int]] {
var numSet = Set<[Int]>()
var numDictionary = [Int: Int]()
let nums = nums.sorted()
for (index, num) in nums.enumerated() {
numDictionary[num] = index
}
for (index1, num1) in nums.enumerated() {
if index1 > 0 && num1 == nums[index1 - 1] {
continue
}
for (index2, num2) in nums[(index1 + 1)...].enumerated() {
let num = 0 - num1 - num2
if let index = numDictionary[num] {
if index > index1 + index2 + 1 {
numSet.insert([num1, num2, num])
}
}
}
}
return Array(numSet)
}
這個代碼雖然最后通過了測試，但是耗時1828ms，只戰(zhàn)勝了14.66%的 Swift 提交記錄。所以代碼優(yōu)化的空間應(yīng)該還是很大的，感覺應(yīng)該還可以通過類似快排的方式進行優(yōu)化。我未來還會針對這個代碼修改，爭取達到主流水準。
