main.go

package main

import "fmt"

// There are N students in a class. Some of them are friends, while some are not.
// Their friendship is transitive in nature. For example, if A is a direct friend of B,
// and B is a direct friend of C, then A is an indirect friend of C. And we defined a
// friend circle is a group of students who are direct or indirect friends.
//
// Given a N*N matrix M representing the friend relationship between students in the class.
// If M[i][j] = 1, then the ith and jth students are direct friends with each other, otherwise not.
// And you have to output the total number of friend circles among all the students.
//
// Example 1:
// Input:
// [[1,1,0],
// [1,1,0],
// [0,0,1]]
// Output: 2
// Explanation:The 0th and 1st students are direct friends, so they are in a friend circle.
// The 2nd student himself is in a friend circle. So return 2.
//
// Example 2:
// Input:
// [[1,1,0],
// [1,1,1],
// [0,1,1]]
// Output: 1
// Explanation:The 0th and 1st students are direct friends, the 1st and 2nd students are direct friends,
// so the 0th and 2nd students are indirect friends. All of them are in the same friend circle, so return 1.
//
// Note:
// N is in range [1,200].
// M[i][i] = 1 for all students.
// If M[i][j] = 1, then M[j][i] = 1.

// 42s, 100%, really that fast?
// quick-find actually, check two student whether is friend is fast
// can't believe quick-find is fastest, must be the testdata is perfect for quick-find
func findCircleNum(M [][]int) int {
	var ret = len(M)
	cycle := make([]int, ret)
	for i := 0; i < ret; i++ {
		cycle[i] = i
	}
	var union = func(s, d int) {
		for i, v := range cycle {
			if v == s {
				cycle[i] = d
			}
		}
	}
	for i := 0; i < len(M); i++ {
		for j := i + 1; j < len(M[i]); j++ {
			if M[i][j] == 1 {
				if cycle[i] != cycle[j] {
					ret--
					union(cycle[j], cycle[i])
				}
			}
		}
	}
	return ret
}

// quick-union
// 56ms
func findCircleNum2(M [][]int) int {
	var ret = len(M)
	cycle := make([]int, ret)
	for i := 0; i < ret; i++ {
		cycle[i] = i
	}
	// find root
	var find = func(s int) int {
		for s != cycle[s] {
			s = cycle[s]
		}
		return s
	}
	for i := 0; i < len(M); i++ {
		for j := i + 1; j < len(M[i]); j++ {
			if M[i][j] == 1 {
				if ri, rj := find(i), find(j); ri != rj {
					ret--
					cycle[rj] = i
				}
			}
		}
	}
	return ret
}

// wighted quick-union
// 52ms
func findCircleNum3(M [][]int) int {
	var ret = len(M)
	cycle := make([]int, ret)
	size := make([]int, ret)
	for i := 0; i < ret; i++ {
		cycle[i] = i
		size[i] = 1
	}

	// find root
	//var find = func(s int) int {
	//	for s != cycle[s] {
	//		s = cycle[s]
	//	}
	//	return s
	//}

	// path compression
	// worse? 62ms?
	var find = func(s int) int {
		t := s
		for s != cycle[s] {
			s = cycle[s]
		}
		cycle[t] = s
		return s
	}

	for i := 0; i < len(M); i++ {
		for j := i + 1; j < len(M[i]); j++ {
			if M[i][j] == 1 {
				if ri, rj := find(i), find(j); ri != rj {
					ret--
					if size[i] < size[j] {
						cycle[ri] = j
						size[j] += size[i]
					} else {
						cycle[rj] = i
						size[i] += size[j]
					}
				}
			}
		}
	}
	return ret
}

func main() {
	friends := [][]int{
		{1, 1, 0},
		{1, 1, 0},
		{0, 0, 1},
	}
	fmt.Println(findCircleNum(friends))
	fmt.Println(findCircleNum2(friends))
	fmt.Println(findCircleNum3(friends))

	friends = [][]int{
		{1, 1, 0},
		{1, 1, 1},
		{0, 1, 1},
	}
	fmt.Println(findCircleNum(friends))
	fmt.Println(findCircleNum2(friends))
	fmt.Println(findCircleNum3(friends))

	friends = [][]int{
		{1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0},
		{1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
		{0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
		{0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0},
		{0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0},
		{0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0},
		{0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0},
		{1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0},
		{0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 0},
		{0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1},
		{0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0},
		{0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0},
		{0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1},
	}
	fmt.Println(findCircleNum(friends))
	fmt.Println(findCircleNum2(friends))
	fmt.Println(findCircleNum3(friends))
}