本题为2024年9月1日字节跳动-秋招机考原题
字节跳动机考的介绍点击这里
小红有一个长度为n的排列p,他想知道p中有多少个i,j对满足:i<j且pi+pj=i+j。
请你帮他算算吧。
输入包含两行。
第一行一个正整数n(1≤n≤200000),表示排列的长度。
第二行n个正整数pi(1≤pi≤n),表示排列p。(保证输入是一个排列。)
输出一行一个整数表示好对的个数。
输入
5
2 1 3 5 4
输出
4
说明
p1+p2=2+1,
p1+p5=2+4,
p2+p4=1+5,
p4+p5=4+5。
共这四对。
这道题要求我们计算在给定的排列 p
中有多少个 (i, j)
对满足条件:
受到哈希3的启发,我们自然想到从左往右枚举每一个数j , 维护一个前缀哈希表H,那么也就是在i∈[1,j−1] 里寻找有多少个i 满足:
pi+pj=i+j而由于在枚举的过程中,我们就已经知道了j , 未知的是i 。那么等式可以转化为:
pi−i=j−pj所以我们
1.在H 里 查询j−pj 的个数 H[j−pj]
2.往H里插入pj−j 。
3.遍历下一个元素
arr = [1, 3, 2, 5]
H = {}
p_1 = 1
, i = 1
H[j - p_j]
,即 H[1 - 1] = H[0]
。H[0]
还没有出现过,所以结果为 0。H
中增加 p_1 - 1 = 1 - 1 = 0
,所以 H = {0: 1}
。p_2 = 3
, i = 2
H[j - p_j] = H[2 - 3] = H[-1]
。H[-1]
还没有出现过,所以结果为 0。H
中增加 p_2 - 2 = 3 - 2 = 1
,所以 H = {0: 1, 1: 1}
。p_3 = 2
, i = 3
H[j - p_j] = H[3 - 2] = H[1]
。H[1]
出现过 1 次,所以结果为 1。H
中增加 p_3 - 3 = 2 - 3 = -1
,所以 H = {0: 1, 1: 1, -1: 1}
。p_4 = 5
, i = 4
H[j - p_j] = H[4 - 5] = H[-1]
。H[-1]
出现过 1 次,所以结果为 1。H
中增加 p_4 - 4 = 5 - 4 = 1
,所以 H = {0: 1, 1: 2, -1: 1}
。最终,我们得到了满足条件的 (i, j)
对的数量为:
j = 2
,没有符合条件的对,结果为 0。j = 3
,找到了 1 个符合条件的对,结果为 1。j = 4
,找到了 1 个符合条件的对,结果为 1。所以,总共有 2 个符合条件的 (i, j)
对。
def count_pairs(n, p):
# 创建一个哈希表,用于记录 p_i - i 出现的次数
diff_count = {}
# 结果变量,记录符合条件的(i, j)对数
result = 0
for j in range(n):
# 计算 j - p_j
diff = (j + 1) - p[j]
# 如果 -diff 已经出现过,那么就说明存在符合条件的 i
if -diff in diff_count:
result += diff_count[-diff]
# 更新 diff_count
if diff in diff_count:
diff_count[diff] += 1
else:
diff_count[diff] = 1
return result
# 输入读取
n = int(input())
p = list(map(int, input().split()))
# 调用函数并输出结果
print(count_pairs(n, p))