923. 3Sum With Multiplicity

Problem

Given an integer array A, and an integer target, return the number of tuples i, j, k such that i < j < k and A[i] + A[j] + A[k] == target.

As the answer can be very large, return it modulo 10^9 + 7.

Example 1:


Input: A = [1,1,2,2,3,3,4,4,5,5], target = 8
Output: 20
Explanation: 
Enumerating by the values (A[i], A[j], A[k]):
(1, 2, 5) occurs 8 times;
(1, 3, 4) occurs 8 times;
(2, 2, 4) occurs 2 times;
(2, 3, 3) occurs 2 times.

Example 2:


Input: A = [1,1,2,2,2,2], target = 5
Output: 12
Explanation: 
A[i] = 1, A[j] = A[k] = 2 occurs 12 times:
We choose one 1 from [1,1] in 2 ways,
and two 2s from [2,2,2,2] in 6 ways.

Note:

3 <= A.length <= 3000
0 <= A[i] <= 100
0 <= target <= 300

Problem Summary

This problem is an upgraded version of Problem 15. Given an array, find the number of combinations of 3 numbers whose sum equals target, with the requirement that i < j < k. The number of combinations of solutions does not need to be deduplicated; the same values with different indices count as different solutions (this is also the difference from Problem 15).

Solution Approach

The general solution for this problem is the same as Problem 15, except that when counting all solution combinations, some knowledge of permutations and combinations is needed. If 3 identical numbers are chosen, calculate C n 3; when choosing 2 identical numbers, calculate C n 2; choosing one number is calculated normally. Finally, add up the number of all solutions.

Code


package leetcode

import (
	"sort"
)

func threeSumMulti(A []int, target int) int {
	mod := 1000000007
	counter := map[int]int{}
	for _, value := range A {
		counter[value]++
	}

	uniqNums := []int{}
	for key := range counter {
		uniqNums = append(uniqNums, key)
	}
	sort.Ints(uniqNums)

	res := 0
	for i := 0; i < len(uniqNums); i++ {
		ni := counter[uniqNums[i]]
		if (uniqNums[i]*3 == target) && counter[uniqNums[i]] >= 3 {
			res += ni * (ni - 1) * (ni - 2) / 6
		}
		for j := i + 1; j < len(uniqNums); j++ {
			nj := counter[uniqNums[j]]
			if (uniqNums[i]*2+uniqNums[j] == target) && counter[uniqNums[i]] > 1 {
				res += ni * (ni - 1) / 2 * nj
			}
			if (uniqNums[j]*2+uniqNums[i] == target) && counter[uniqNums[j]] > 1 {
				res += nj * (nj - 1) / 2 * ni
			}
			c := target - uniqNums[i] - uniqNums[j]
			if c > uniqNums[j] && counter[c] > 0 {
				res += ni * nj * counter[c]
			}
		}
	}
	return res % mod
}