[LeetCode] 1995. Count Special Quadruplets

轉自LeetCode

Given a 0-indexed integer array nums, return the number of distinct quadruplets (a, b, c, d) such that:

• nums[a] + nums[b] + nums[c] == nums[d], and
• a < b < c < d

 

Example 1:

Input: nums = [1,2,3,6]
Output: 1
Explanation: The only quadruplet that satisfies the requirement is (0, 1, 2, 3) because 1 + 2 + 3 == 6.

Example 2:

Input: nums = [3,3,6,4,5]
Output: 0
Explanation: There are no such quadruplets in [3,3,6,4,5].

Example 3:

Input: nums = [1,1,1,3,5]
Output: 4
Explanation: The 4 quadruplets that satisfy the requirement are:
- (0, 1, 2, 3): 1 + 1 + 1 == 3
- (0, 1, 3, 4): 1 + 1 + 3 == 5
- (0, 2, 3, 4): 1 + 1 + 3 == 5
- (1, 2, 3, 4): 1 + 1 + 3 == 5

 

Constraints:

• 4 <= nums.length <= 50
• 1 <= nums[i] <= 100

Solution

這題想到的是 prefixSum 的方向，也是得對數學式做個手腳

nums[a] + nums[b] + nums[c] = nums[d]

變成

nums[a] + nums[b]  = nums[d] - nums[c]

那因為題目也有規定 a < b < c < d

所以可以專心掃過 a 和 b 的組合，然後用 nums[a] + nums[b] 來當 key 找出有多少組合

需要用一個 map 來記錄 nums[d] - nums[c] 有多少個 (remember: nums[a] + nums[b]  = nums[d] - nums[c] )

由後往前算，這樣 map 才能被正確更新

用 1, 1, 1, 3, 5 來舉例

初始值：

nums[d] - nums[c] = nums[4] - nums[3] = 5 - 3 = 2

map[nums[d] - nums[c]] = map[2] = 1

1st round:

1, 1, 1, 3, 5 

   a  b

a 往前移動，(a,b) 有可能的組合為 (1,1) 一種，然後 nums[a] + nums[b] = 2

ans += if (map[2] == null) 0 else map[2]!!

這時也要更新 map 給下一輪使用，範圍是 b+1 到 最後一個數

1, 1, 1, 3, 5 

      b     d

map[nums[4] - nums[2]] + 1= map[4] + 1

map[nums[3] - nums[2]] + 1= map[2] + 1

更新的邏輯是增加 count，如果 key 存在，則把 count + 1，不存在，就設為 1

2nd round

1, 1, 1, 3, 5 

a  b     

a 往前移動，(a,b) 有可能的組合為 (1,1) 一種，然後 nums[a] + nums[b] = 2

ans += if (map[2] == null) 0 else map[2]!!

更新 map

map[nums[4] - nums[1]] + 1= map[4] + 1

map[nums[3] - nums[1]] + 1= map[2] + 1

map[nums[2] - nums[1]] + 1= map[0] + 1

最後就可以找到答案

kotlin(參考解答)

	class Solution {
	fun countQuadruplets(nums: IntArray): Int {
	val map = mutableMapOf<Int,Int>()
	var ans = 0
	map[nums[nums.lastIndex] - nums[nums.lastIndex-1]] = 1
	// nums[a] + nums[b] = nums[d] - nums[c]
	// a < b < c < d
	for(b in nums.lastIndex-2 downTo 1) {
	for(a in b-1 downTo 0) {
	ans += map[nums[a] + nums[b]] ?: 0
	}
	for(d in nums.lastIndex downTo b+1) {
	val key = nums[d] - nums[b]
	map[key] = map[key]?.let{it+1} ?: 1
	}
	}
	return ans
	}
	}

view raw count_special_quadruplets.kt hosted with ❤ by GitHub