美团2024春 小美的朋友关系 Golang题解

小美的朋友关系

小美认为，在人际交往中，但是随着时间的流逝，朋友的关系也是会慢慢变淡的，最终朋友关系就淡忘了。
现在初始有一些朋友关系，存在一些事件会导致两个人淡忘了他们的朋友关系。小美想知道某一时刻中，某两人是否可以通过朋友介绍互相认识？
事件共有 2 种：
1 u v：代表编号 u 的人和编号 v 的人淡忘了他们的朋友关系。
2 u v：代表小美查询编号 u 的人和编号 v 的人是否能通过朋友介绍互相认识。

注：介绍可以有多层，比如 2 号把 1 号介绍给 3 号，然后 3 号再把 1 号介绍给 4 号，这样 1 号和 4 号就认识了。

时间限制：C/C++ 1秒，其他语言2秒

空间限制：C/C++ 256M，其他语言512M

第一行输入三个正整数，代表总人数，初始的朋友关系数量，发生的事件数量。
接下来的行，每行输入两个正整数，代表初始编号的人和编号的人是朋友关系。
接下来的行，每行输入三个正整数，含义如题目描述所述。
1 <= n <= 10^9
1 <= m,q <= 10^5
1 <= u,v <= n
1 <= op <= 2
保证至少存在一次查询操作。

先说思路：本题采用数据离散化+倒序并查集即可AC

因为并查集不支持删除关系，所以我们需要先将所有删除和查询操作记录下来，然后从后往前遍历所有的操作，此时需要将删除操作视为合并并查集的操作。这一过程类似于时光回溯，你可以想象一下

package main

import (
	"bufio"
	"fmt"
	"os"
)

var (
	root []int
	rank []int
)

func main() {
	reader := bufio.NewReaderSize(os.Stdin, 1<<20)
	_, m, q := readInt(reader), readInt(reader), readInt(reader)
	nodeID := make(map[int]int)
	idCnt := 1

	// 记录已有友情
	relationship := make(map[int64]bool, m)
	for i := 0; i < m; i++ {
		a, b := readInt(reader), readInt(reader)
		aID := getID(a, nodeID, &idCnt)
		bID := getID(b, nodeID, &idCnt)
		relationship[encode(aID, bID)] = true
	}

	// 记录是否是原来需要的边
	// 试想一下，因为我们要从后往前遍历options数组，将删除操作转化为增加操作
	// 如果两个人本来就不认识，但现在有一个操作是删除他俩之间的关系，那我们就会将这个删除操作变成增加操作，也就是一番操作后本来就不该相识的俩人突然认出了彼此。这是不该出现的情况
	// 所以这里需要一个数组记录俩人是否原本就相识，确保后续将删除操作转为增加操作时不会误增加关系
	need := make(map[int64]bool, m)

	// 记录所有操作
	options := make([][]int, 0)
	for i := 0; i < q; i++ {
		o, a, b := readInt(reader), readInt(reader), readInt(reader)
		aID := getID(a, nodeID, &idCnt)
		bID := getID(b, nodeID, &idCnt)

		options = append(options, []int{o, aID, bID})
		if o == 1 {
			key := encode(aID, bID)
			// 判断俩人是否原本就认识
			if _, ok := relationship[key]; ok {
				delete(relationship, key)
				need[key] = true
			}
		}
	}

	initArr(idCnt)

	// 使用删除关系后的结果创建并查集
	for k := range relationship {
		a, b := decode(k)
		union(a, b)
	}

	ans := make([]string, 0)
	for i := len(options) - 1; i >= 0; i-- {
		o, a, b := options[i][0], options[i][1], options[i][2]
		if o == 1 {
			key := encode(a, b)
			// 如果俩人本来就相识，才能将删除转为新增关系
			if _, ok := need[key]; ok {
				union(a, b)
			}
		} else {
			if find(a) == find(b) {
				ans = append(ans, "Yes")
			} else {
				ans = append(ans, "No")
			}
		}
	}

	for i := len(ans) - 1; i >= 0; i-- {
		fmt.Println(ans[i])
	}
}

func getID(x int, nodeID map[int]int, idCnt *int) int {
	if id, ok := nodeID[x]; ok {
		return id
	}
	nodeID[x] = *idCnt
	*idCnt++
	return nodeID[x]
}

func initArr(n int) {
	root = make([]int, n+1)
	rank = make([]int, n+1)
	for i := 1; i <= n; i++ {
		root[i] = i // 每个节点初始状态下的根节点都是自身
		rank[i] = 1 // 每个节点属于的集合在初始状态下内部只有一个节点
	}
}

// 如果该节点的根节点是自身，则返回该节点
// 否则递归查找（顺便压缩路径）
func find(x int) int {
	if x == root[x] {
		return x
	} else {
		root[x] = find(root[x])
		return root[x]
	}
}

// 按秩合并，可以将集合高度或是集合中元素数量作为秩
func union(i, j int) {
	x := find(i)
	y := find(j)
	if x == y {
		return
	}
	if rank[x] < rank[y] {
		root[x] = y
	} else {
		root[y] = x
	}

	if rank[x] == rank[y] {
		rank[x]++
	}
}

func readInt(r *bufio.Reader) int {
	b, _ := r.ReadByte()
	for b < '0' || b > '9' {
		b, _ = r.ReadByte()
	}
	num := 0
	for '0' <= b && b <= '9' {
		num = num*10 + int(b-'0')
		b, _ = r.ReadByte()
	}
	return num
}

// string哈希代价更高，这里换成int64可以再偷点时间
// 高32位存a，低32位存b
func encode(a, b int) int64 {
	// 保证a和b中，小的数字排前面，这样做是为了避免路径重复
	if a > b {
		a, b = b, a
	}
	return int64(a)<<32 | int64(b)
}

func decode(key int64) (int, int) {
	a := int(key >> 32)
	b := int(key & 0xffffffff)
	return a, b
}

这道题卡了我很久，要么过不去用例要么超时爆内存，学习了一番优化思路终于AC了