Completing a Tree

 

http://rosalind.info/problems/tree/

ROSALIND | Completing a Tree

adjacency list를 입력으로 주고 그 리스트들을 하나의 tree로 만들려면 몇개의 최소 edge가 필요한지 묻는 문제이다. cluster별로 분류한 후 cluster의 갯수 - 1을 구하면 된다.

 

from functools import reduce

with open('rosalind_tree.txt') as f:
    nodes, *edges = f.read().strip().split('\n')

segments = []

adj_nodes = reduce(lambda a, b: set(a).union(set(b)), map(lambda x: x.split(' '), edges))
for node in range(1, int(nodes) + 1):
    if str(node) not in adj_nodes:
        segments.append(set([str(node)]))

for edge in edges:
    new_cluster = True
    e1, e2 = edge.split(' ')
    for segment in segments:
        if e1 in segment or e2 in segment:
            segment.add(e1)
            segment.add(e2)
            new_cluster = False
    if new_cluster is True:
        segments.append(set([e1, e2]))

def merge_dup(segments):
    new_segments = segments[0:]
    for seg1 in segments:
        for seg2 in segments:
            if seg1 != seg2 and len(seg1.intersection(seg2)) > 0:
                new_seg = seg1.union(seg2)
                new_segments.remove(seg1)
                new_segments.remove(seg2)
                new_segments.append(new_seg)
                return new_segments
    return segments

prev_seg_len = 0
while prev_seg_len != len(segments):
    prev_seg_len = len(segments)
    segments = merge_dup(segments)

print(len(segments) - 1)

 

for node in range(1, int(nodes) + 1):
    if str(node) not in adj_nodes:
        segments.append(set([str(node)]))

문제에선 입력의 첫번째줄에 노드의 갯수를 주는데 adjacency list에 없는 노드들도 있으므로 해당 노드들은 미리 clusters(segment)에 추가해준다.

 

for edge in edges:
    new_cluster = True
    e1, e2 = edge.split(' ')
    for segment in segments:
        if e1 in segment or e2 in segment:
            segment.add(e1)
            segment.add(e2)
            new_cluster = False
    if new_cluster is True:
        segments.append(set([e1, e2]))

그리고 adjecency list를 순회하며 segment에 이미 해당 edge를 구성하는 node가 있으면 해당 segment에 추가하고 없으면 새로운 segment를 생성하여 segments에 추가한다.

 

def merge_dup(segments):
    new_segments = segments[0:]
    for seg1 in segments:
        for seg2 in segments:
            if seg1 != seg2 and len(seg1.intersection(seg2)) > 0:
                new_seg = seg1.union(seg2)
                new_segments.remove(seg1)
                new_segments.remove(seg2)
                new_segments.append(new_seg)
                return new_segments
    return segments

이 부분이 중요한데 adjecency list를 순회하면서 segments에 추가하면 adjecency list가 정렬되어있다면 별 문제가 안되지만 정렬되어있지 않다면 같은 node를 포함한, 즉 같은 segment가 다른 segment로 추가될 수 있다. 예를들어

1 2
3 4
2 3

이 입력으로 들어온다면

[{1,2}]
-> [{1,2}, {3,4}]
-> [{1,2,3}, {3,4}]

이런식으로 segments가 생성이 될것이다. 그러므로 merge_dup 함수를 만들어 같은 segment지만 분리된것들을 합쳐주었다.