字节跳动2018.9.9笔试 组织优化

题目描述

给定一个M*M的二维数组，每个值1的元素代表一个团队。如果两个团队在上下或左右的方向上相邻，说明2个团队有联系，就要把他们合并到一个部门；没有联系的，就放在不同部门。
判断输入中，有多少个部门。
输入：
第一行，一个整数，M
后面M行，每行M个整数（取值0或者1）
输出：
一个整数，部门数
样例输入：
1 0 0 1 1
1 0 0 1 1
0 0 1 0 0
0 0 1 0 0
0 0 1 0 0
样例输出：
3
提示：
左上角有一个部门，右上角有一个部门，下面有一个部门。

思路

把每个部门都看成一个子树，使用parent数组来存每个节点的父亲，子树里根节点的parent是自身，其他节点直接或者间接指向根节点（间接是指递归地找父节点能找到根节点）。
初始化时，每个节点就是一颗子树。即每个节点的parent就是自身。
最终，查看数据中所有为1的节点中，父节点是自身的数量，返回这个数量。因为每棵树中只有根节点的parent是自身，所以通过这个标准来判断根节点的数量即树的数量。
本题思路与无向图检测环里面的思想很像，读者可以看看，受点启发。

代码

#提交代码时，取消以下注释
##m = eval(input())
##li = []
##for i in range(m):
##    a = list(map(int,input().split()))
##    li.append(a.copy())

#提交代码时，删除以下代码
m = 6
li = [[1,0,0,1,1,1],
      [1,0,0,1,1,1],
      [0,0,1,0,0,0],
      [1,1,1,0,1,0],
      [1,0,1,1,1,0],
      [1,0,1,1,1,0],
    ]

parent = []#把原数组平铺开来变成一维数组，元素个数是一样的
for i in range(m**2):
    parent.append(i)
#初始每个节点的父亲都是自己，代表自己就是一个子树
#注意初始值不能是-1，因为这样就不方便findroot函数返回值比较

def findroot(i):
    #当前节点所在子树的根节点
    if(parent[i] == i):
        return i
    return findroot(parent[i])

#判断点与点的每条边
#先判断横边，ij确定横边左边的那个点
for i in range(m):
    for j in range(m-1):
        if li[i][j] == li[i][j+1] == 1:
            one = i*m+j
            two = i*m+j+1
            oneroot = findroot(one)
            tworoot = findroot(two)
            if(oneroot != tworoot):#不一样说明分别在两颗不同的子树
                #合并两颗子树
                parent[tworoot] = oneroot
                #本来应该parent[tworoot] = one，这里压缩路径，递归次数会减小

#再判断竖边，ij确定竖边上边的那个点
#与上面的不同只是one，two
for j in range(m):
    for i in range(m-1):
        if li[i][j] == li[i+1][j] == 1:
            one = i*m+j
            two = (i+1)*m+j
            oneroot = findroot(one)
            tworoot = findroot(two)
            if(oneroot != tworoot):
                #合并两颗子树
                parent[tworoot] = oneroot

count = 0
for i in range(m):
    for j in range(m):
        flat = i*m + j#ij对应在parent中的索引为flat
        if (li[i][j] == 1) and (parent[flat] == flat):
            count += 1
print(count)

代码中数据与原样例输入不一样，运行结果为3。