USACO SuperBull / Kruskal's Algorithm

USACO 2015 February Silver Problem 3¶

Superbull, solved with Kruskal's Algorithm using Union-Find¶

http://usaco.org/index.php?page=viewproblem2&cpid=531

In [1]:

# SAMPLE INPUT:
#
#
# 4
# 3
# 6
# 9
# 10

In [2]:

fin = open ('superbull.in', 'r')
fout = open ('superbull.out', 'w')
N = int(fin.readline().strip())
ID = []
PARENT = [0]*N
for i in range(N):
    ident = int(fin.readline().strip())
    ID.append(ident)
    PARENT[i]=i
  
list1 = []
print("team1","team2","score")
for t1 in range(N):
    for t2 in range(t1+1,N):
        print(t1,t2,-1*(ID[t1]^ID[t2]))
        list1.append( (-1*(ID[t1]^ID[t2]),t1,t2) )

print("\nedges = ",list1)

team1 team2 score
0 1 -5
0 2 -10
0 3 -9
1 2 -15
1 3 -12
2 3 -3

edges =  [(-5, 0, 1), (-10, 0, 2), (-9, 0, 3), (-15, 1, 2), (-12, 1, 3), (-3, 2, 3)]

In [3]:

# convert edge list to a heap
import heapq
heapq.heapify(list1)
print(list1)

[(-15, 1, 2), (-12, 1, 3), (-9, 0, 3), (-10, 0, 2), (-5, 0, 1), (-3, 2, 3)]

In [4]:

def find(x):
    while x!=PARENT[x]:
        x=PARENT[x]
    return x

In [5]:

def union(x,y):
    a = find(x)
    b = find(y)
    PARENT[a]=b
    return

In [6]:

# replace calls to find() with this call if you want to do path compression
def find_with_path_compression(x):
    if x!=PARENT[x]:
        PARENT[x]=find_with_path_compression(PARENT[x])
    return PARENT[x]

In [7]:

result = 0
while (len(list1)>0):
    edge = heapq.heappop(list1)
    if find(edge[1])!=find(edge[2]):
        print("adding edge:",edge)
        result+=edge[0]
        union(edge[1],edge[2])

result *= -1
print("\nMax weight spanning tree:",result)

adding edge: (-15, 1, 2)
adding edge: (-12, 1, 3)
adding edge: (-10, 0, 2)

Max weight spanning tree: 37

In [ ]:

fout.write(str(result)+"\n")
fout.close()