532. 数组中的 k-diff 数对

题目描述

输入整数数组 nums 和非负整数 k。若两个不同下标 i != j 上的值满足 |nums[i] - nums[j]| == k，就称这两个值构成一个 k-diff 数对。

统计有多少种不同的 k-diff 数对并返回。判断是否相同只看数值：值组合一样的数对只记一次，与它们出现的下标和次数无关，(a, b) 与 (b, a) 视为同一对。

当 k = 0 时，只有在数组中出现至少两次的值才能和自己配成一对。

示例 1

输入：nums = [3,1,4,1,5], k = 2
输出：2

解释：满足差为 2 的值组合有 (1,3) 和 (3,5)，虽然 1 出现了两次，(1,3) 仍只算一次。

示例 2

输入：nums = [1,2,3,4,5], k = 1
输出：4

解释：相邻整数 (1,2)、(2,3)、(3,4)、(4,5) 共 4 对。

示例 3

输入：nums = [1,3,1,5,4], k = 0
输出：1

解释：只有 1 出现了两次，能组成 (1,1) 这一对。

数据范围

1 <= nums.length <= 10^4
-10^7 <= nums[i] <= 10^7
0 <= k <= 10^7

解题思路：排序 + 二分查找

排序后，每对只在较小的数 x 处计一次：对每个不重复的 x，只在它右边 [i+1, n) 二分查找 x + k。跳过重复的 x 保证同一对不重计；从 i+1 开始搜保证 i != j，k=0 时也只有重复值才能配对。

为什么正确

处理到第 i 个数时，ans 等于较小端严格小于 nums[i] 的不同数对个数。

复杂度分析

时间 O(n log n)（排序 + n 次二分），额外空间 O(1)（不计排序）。

易错点

k=0 时如果从 i 本身开始搜，会把自己和自己配成一对；必须从 i+1 开始。不跳过重复的 x 则会把 (1,3) 算两次。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

class Solution:
    def findPairs(self, nums: List[int], k: int) -> int:
        nums.sort()
        n = len(nums)
        ans = 0
        for i in range(n):
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            j = bisect_left(nums, nums[i] + k, i + 1)
            if j < n and nums[j] == nums[i] + k:
                ans += 1
        return ans

class Solution {
    public int findPairs(int[] nums, int k) {
        java.util.Arrays.sort(nums);
        int n = nums.length;
        int ans = 0;
        for (int i = 0; i < n; i++) {
            if (i > 0 && nums[i] == nums[i - 1])
                continue;
            int lo = i + 1, hi = n; while (lo < hi) { int m = (lo + hi) >>> 1; if (nums[m] < nums[i] + k) lo = m + 1; else hi = m; } int j = lo;
            if (j < n && nums[j] == nums[i] + k)
                ans++;
        }
        return ans;
    }
}

class Solution {
public:
    int findPairs(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        int n = nums.size();
        int ans = 0;
        for (int i = 0; i < n; i++) {
            if (i > 0 && nums[i] == nums[i - 1])
                continue;
            int j = lower_bound(nums.begin() + i + 1, nums.end(), nums[i] + k) - nums.begin();
            if (j < n && nums[j] == nums[i] + k)
                ans++;
        }
        return ans;
    }
};

func findPairs(nums []int, k int) int {
	sort.Ints(nums)
	n := len(nums)
	ans := 0
	for i := 0; i < n; i++ {
		if i > 0 && nums[i] == nums[i-1] {
			continue
		}
		j := i + 1 + sort.SearchInts(nums[i+1:], nums[i]+k)
		if j < n && nums[j] == nums[i]+k {
			ans++
		}
	}
	return ans
}

var findPairs = function(nums, k) {
    nums.sort((a, b) => a - b);
    const n = nums.length;
    let ans = 0;
    for (let i = 0; i < n; i++) {
        if (i > 0 && nums[i] === nums[i - 1])
            continue;
        let lo = i + 1, hi = n; while (lo < hi) { const m = (lo + hi) >> 1; if (nums[m] < nums[i] + k) lo = m + 1; else hi = m; } const j = lo;
        if (j < n && nums[j] === nums[i] + k)
            ans++;
    }
    return ans;
};

function findPairs(nums: number[], k: number): number {
    nums.sort((a, b) => a - b);
    const n = nums.length;
    let ans = 0;
    for (let i = 0; i < n; i++) {
        if (i > 0 && nums[i] === nums[i - 1])
            continue;
        let lo = i + 1, hi = n; while (lo < hi) { const m = (lo + hi) >> 1; if (nums[m] < nums[i] + k) lo = m + 1; else hi = m; } const j = lo;
        if (j < n && nums[j] === nums[i] + k)
            ans++;
    }
    return ans;
}

载入编辑器，自己跑一遍 →

其他解法

三种解法都只从较小值出发找 x + k 来去重；哈希表最快但占空间，排序后用二分或双指针都只需排序的空间，双指针的扫描部分更省时。

解法 2：哈希表计数时间 O(n)，空间 O(n)

先用哈希表统计每个数出现的次数。k = 0 时，出现次数不少于 2 的数各贡献一个数对。k > 0 时，对每个不同的数 x，检查 x + k 是否在表中。只向较大一侧找，保证每个数对只数一次。

不需要排序，时间降到 O(n)，但要 O(n) 额外空间。k = 0 要单独处理，是最容易出错的地方。

class Solution:
    def findPairs(self, nums: List[int], k: int) -> int:
        cnt = {}
        for x in nums:
            cnt[x] = cnt.get(x, 0) + 1
        ans = 0
        for x, c in cnt.items():
            if k == 0:
                if c > 1:
                    ans += 1
            elif x + k in cnt:
                ans += 1
        return ans

class Solution {
    public int findPairs(int[] nums, int k) {
        Map<Integer, Integer> cnt = new HashMap<>();
        for (int x : nums) cnt.merge(x, 1, Integer::sum);
        int ans = 0;
        for (Map.Entry<Integer, Integer> e : cnt.entrySet()) {
            int x = e.getKey(), c = e.getValue();
            if (k == 0) {
                if (c > 1) ans++;
            } else if (cnt.containsKey(x + k)) {
                ans++;
            }
        }
        return ans;
    }
}

class Solution {
public:
    int findPairs(vector<int>& nums, int k) {
        unordered_map<int, int> cnt;
        for (int x : nums) cnt[x]++;
        int ans = 0;
        for (auto& [x, c] : cnt) {
            if (k == 0) {
                if (c > 1) ans++;
            } else if (cnt.count(x + k)) {
                ans++;
            }
        }
        return ans;
    }
};

func findPairs(nums []int, k int) int {
	cnt := map[int]int{}
	for _, x := range nums {
		cnt[x]++
	}
	ans := 0
	for x, c := range cnt {
		if k == 0 {
			if c > 1 {
				ans++
			}
		} else if _, ok := cnt[x+k]; ok {
			ans++
		}
	}
	return ans
}

var findPairs = function(nums, k) {
    const cnt = new Map();
    for (const x of nums) cnt.set(x, (cnt.get(x) || 0) + 1);
    let ans = 0;
    for (const [x, c] of cnt) {
        if (k === 0) {
            if (c > 1) ans++;
        } else if (cnt.has(x + k)) {
            ans++;
        }
    }
    return ans;
};

function findPairs(nums: number[], k: number): number {
    const cnt = new Map<number, number>();
    for (const x of nums) cnt.set(x, (cnt.get(x) || 0) + 1);
    let ans = 0;
    for (const [x, c] of cnt) {
        if (k === 0) {
            if (c > 1) ans++;
        } else if (cnt.has(x + k)) {
            ans++;
        }
    }
    return ans;
}

解法 3：排序 + 双指针时间 O(n log n)，空间 O(log n)

排序后，i 枚举每个不同的数作为较小值。j 从 i + 1 之后出发，只向右移动，直到 nums[j] - nums[i] >= k。差值恰好等于 k 时计数加一。nums[i] 递增，所以 j 不需要回退。

排序之后的扫描是 O(n)，比逐个二分少一个 log 因子。代码要注意 j 至少在 i 右边。

class Solution:
    def findPairs(self, nums: List[int], k: int) -> int:
        nums.sort()
        n = len(nums)
        ans = j = 0
        for i in range(n):
            if i > 0 and nums[i] == nums[i - 1]:
                continue
            j = max(j, i + 1)
            while j < n and nums[j] - nums[i] < k:
                j += 1
            if j < n and nums[j] - nums[i] == k:
                ans += 1
        return ans

class Solution {
    public int findPairs(int[] nums, int k) {
        Arrays.sort(nums);
        int n = nums.length, ans = 0, j = 0;
        for (int i = 0; i < n; i++) {
            if (i > 0 && nums[i] == nums[i - 1]) continue;
            j = Math.max(j, i + 1);
            while (j < n && nums[j] - nums[i] < k) j++;
            if (j < n && nums[j] - nums[i] == k) ans++;
        }
        return ans;
    }
}

class Solution {
public:
    int findPairs(vector<int>& nums, int k) {
        sort(nums.begin(), nums.end());
        int n = nums.size(), ans = 0, j = 0;
        for (int i = 0; i < n; i++) {
            if (i > 0 && nums[i] == nums[i - 1]) continue;
            j = max(j, i + 1);
            while (j < n && nums[j] - nums[i] < k) j++;
            if (j < n && nums[j] - nums[i] == k) ans++;
        }
        return ans;
    }
};

func findPairs(nums []int, k int) int {
	sort.Ints(nums)
	n, ans, j := len(nums), 0, 0
	for i := 0; i < n; i++ {
		if i > 0 && nums[i] == nums[i-1] {
			continue
		}
		j = max(j, i+1)
		for j < n && nums[j]-nums[i] < k {
			j++
		}
		if j < n && nums[j]-nums[i] == k {
			ans++
		}
	}
	return ans
}

var findPairs = function(nums, k) {
    nums.sort((a, b) => a - b);
    const n = nums.length;
    let ans = 0, j = 0;
    for (let i = 0; i < n; i++) {
        if (i > 0 && nums[i] === nums[i - 1]) continue;
        j = Math.max(j, i + 1);
        while (j < n && nums[j] - nums[i] < k) j++;
        if (j < n && nums[j] - nums[i] === k) ans++;
    }
    return ans;
};

function findPairs(nums: number[], k: number): number {
    nums.sort((a, b) => a - b);
    const n = nums.length;
    let ans = 0, j = 0;
    for (let i = 0; i < n; i++) {
        if (i > 0 && nums[i] === nums[i - 1]) continue;
        j = Math.max(j, i + 1);
        while (j < n && nums[j] - nums[i] < k) j++;
        if (j < n && nums[j] - nums[i] === k) ans++;
    }
    return ans;
}

二分查找专题的其他题目

查看二分查找全部题目 →