【蓝桥杯】省赛：缴纳过路费（并查集）

思路

这题很有意思，路线以最贵的那一段收费w为标准，而只要w在区间[l,r]之间即可。
所以城市对之间必须至少有一段路价格不小于l，至于大于r（超费）的输入，直接不理他（不加入图）即可。
最后求的是城市对的数量，可以用并查集，把路段价格小于R满足的城市加入到集合A中，再把所有路段价格小于l的城市再存入另一个集合B，最后两两组合配对，从集合A中所有配对数量减去B中配对数量。所以额外需要记录根节点的size
$$pair=(n-1)+(n-2)+...+1=n\ast (n-1)//2$$
最终
$$ans=pairA-pairB$$

code

因为要创建两个并查集，这里把并查集封装成对象更好，但是不熟悉我就用参数a来区分两个集合了，稍微麻烦。

import os
import sys

def find(x,a):
  if a==1:
    if father1[x]==x:return x
    father1[x] = find(father1[x],a)
    return father1[x]
  if a==2:
    if father2[x]==x:return x
    father2[x] = find(father2[x],a)
    return father2[x]

def merge(x,y,a):
  if a==1:
    fx = find(x,a)
    fy = find(y,a)
    if fx != fy:
      father1[fy] = fx
      size1[fx] += size1[fy]
      size1[fy] = 1
  if a==2:
    fx = find(x,a)
    fy = find(y,a)
    if fx != fy:
      father2[fy] = fx
      size2[fx] += size2[fy]
      size2[fy] = 1

def cal(x):
  return x*(x-1)//2

n, m, L, R = map(int, input().split())
road = []

father1 = list(range(n+1))
size1 = [1 for i in range(n+1)]
father2 = list(range(n+1))
size2 = [1 for i in range(n+1)]
for i in range(m):
  u, v, w = map(int, input().split())
  if w <= R: merge(u,v,1)
  if w < L: merge(u,v,2)

ans = 0
for i in range(1,n+1):
  if i==find(i,1):
    ans += cal(size1[i])
  if i==find(i,2):
    ans -= cal(size2[i])
print(ans)