1534. 统计好三元组

Problem#

给你一个整数数组 arr，以及三个整数 a、b、c。

请统计其中好三元组的数量。

如果三元组 (arr[i], arr[j], arr[k]) 满足下列全部条件，则认为它是一个 好三元组：

0 <= i < j < k < arr.length
|arr[i] - arr[j]| <= a
|arr[j] - arr[k]| <= b
|arr[i] - arr[k]| <= c

https://leetcode.cn/problems/count-good-triplets/description/

示例 1：

输入：arr = [3,0,1,1,9,7], a = 7, b = 2, c = 3 输出：4 解释：一共有 4 个好三元组：(3,0,1), (3,0,1), (3,1,1), (0,1,1)

示例 2：

输入：arr = [1,1,2,2,3], a = 0, b = 0, c = 1 输出：0

约束：

3 <= arr.length <= 100
0 <= arr[i] <= 1000
0 <= a, b, c <= 1000

思路#

三重循环遍历所有 i < j < k 组合，逐一判断条件。

1
class Solution:
2
    def countGoodTriplets(self, arr: list[int], a: int, b: int, c: int) -> int:
3
        n = len(arr)
4
        count = 0
5
        for i in range(n - 2):
6
            for j in range(i + 1, n - 1):
7
                if abs(arr[i] - arr[j]) > a:
8
                    continue
9
                for k in range(j + 1, n):
10
                    if abs(arr[j] - arr[k]) <= b and abs(arr[i] - arr[k]) <= c:
11
                        count += 1
12
        return count

1
import unittest
2

3

4
class TestCountGoodTriplets(unittest.TestCase):
5
    def setUp(self):
6
        self.solution = Solution()
7

8
    def test_example1(self):
9
        self.assertEqual(self.solution.countGoodTriplets([3, 0, 1, 1, 9, 7], 7, 2, 3), 4)
10

11
    def test_example2(self):
12
        self.assertEqual(self.solution.countGoodTriplets([1, 1, 2, 2, 3], 0, 0, 1), 0)
13

14
    def test_all_valid(self):
15
        self.assertEqual(self.solution.countGoodTriplets([1, 1, 1], 0, 0, 0), 1)
16

17

18
if __name__ == '__main__':
19
    unittest.main()