Python中的对角线问题

题目链接：蓝桥杯2024年第十五届省赛真题-连连看 - C语言网

做法一：基于元素

n,m=map(int,input().split())
mix=[list(map(int,input().split())) for i in range(n)]
N=int(2e3 + 5)
st1=[[0]*N for _ in range(N)]                                                                                   
st2=[[0]*N for _ in range(N)]
for i in range(n):
    for j in range(m):
        st1[i-j+1000][mix[i][j]]+=1
        st2[i+j][mix[i][j]]+=1
ans=0
for i in range(n):
    for j in range(m):
        ans+=st1[i-j+1000][mix[i][j]]-1
        ans+=st2[i+j][mix[i][j]]-1
print(ans)

1. 预处理：建立对角线的统计表

代码通过两个二维数组 st1 和 st2 分别记录两种方向的对角线特征：

• 正对角线（左上到右下）：用 i-j 标识每条对角线。为避免负数索引，统一加上偏移量 1000（例如，当 i=0, j=5 时，索引为 0-5+1000=995）。
• 反对角线（右上到左下）：用 i+j 直接标识，无需偏移（例如，i=2, j=3 对应索引 5）。

遍历矩阵时，每个元素的值会被记录到对应的对角线统计表中。例如，若 mix[i][j]=5，则 st1[i-j+1000] 和 st2[i+j] 的计数各加 1。

---

2. 计算相同值的元素对

第二次遍历矩阵时，对于每个元素 mix[i][j]，代码通过以下方式统计有效对数：

• 正对角线上的对数：st1[i-j+1000][mix[i][j]] - 1。减 1 是为了排除自身的一次计数，仅统计同一对角线上其他相同值的元素。
• 反对角线上的对数：同理，st2[i+j][mix[i][j]] - 1。

最终，所有元素的对数累加得到总和 ans。例如，若某对角线有 3 个值为 5 的元素，则每个元素贡献 3-1=2 次，总贡献为 3×2=6。

 做法二：基于对角线

n,m=map(int,input().split())
mix=[]
for _ in range(n):
    mix.append(list(map(int,input().split())))
N=int(2e3 + 5)
st1={}
st2={}
ans=0

for i in range(n):
    for j in range(m):
        #添加对角线
        if i-j+1000 not in st1:
            st1[i-j+1000]={}
        if i+j not in st2:
            st2[i+j]={}
        #添加元素索引
        if mix[i][j] not in st1[i-j+1000]:
            ans+=0
            st1[i-j+1000][mix[i][j]]=0
        else:
            ans+=st1[i-j+1000][mix[i][j]]*2
        st1[i-j+1000][mix[i][j]]+=1
        if mix[i][j] not in st2[i+j]:
            ans+=0
            st2[i+j][mix[i][j]]=0
        else:
            ans+=st2[i+j][mix[i][j]]*2
        st2[i+j][mix[i][j]]+=1
print(ans)

1. 预处理对角线与反对角线

• 正对角线标识：通过 i-j+1000 标识每条正对角线，+1000 用于避免负索引（例如 i=0, j=5 对应索引 995）。
• 反对角线标识：通过 i+j 直接标识反对角线（例如 i=2, j=3 对应索引 5）。
• 动态字典存储：使用嵌套字典 st1 和 st2 分别记录每条对角线上各数值的出现次数。字典结构相比固定数组更灵活，节省内存空间。

2. 动态统计元素对

遍历矩阵时，对每个元素 mix[i][j] 执行以下操作：

1. 检查字典初始化：若当前对角线或反对角线未记录过该数值，则初始化对应键值。
2. 累加元素对数量：
   • 若当前值已存在于当前对角线，则 ans += st1[i-j+1000][mix[i][j]] * 2；
   • 同理处理反对角线 st2[i+j][mix[i][j]]。
3. 更新计数器：将当前值的计数加1。

示例：假设某对角线已有2个值为5的元素，当第三个5出现时，ans 会增加 2*2=4（即每个已有元素与新元素形成2对，共2×2=4次累加）。

---

更多优质学习内容——微信公众号：“以学会友ing”