SS'
SS'

Reputation: 857

Construct tree given root

ex. [13,11], [13,8], [6, 8], [3, 6] where the root is 1

Looking for a Pythonic way to construct a tree into a dictionary such that we then have {13: [11, 8], 11: [], 3: [], 6: [3], 8: [6]}

Since I know the root, the way I would approach is to loop through entries with 13, connect those, then treating those entries as root, so and so forth.

Note there is no ordering. To distinguish between, I know the root (some number) and can determine from there.

Is there a better way?

Upvotes: 3

Views: 661

Answers (5)

Ajax1234
Ajax1234

Reputation: 71451

You can use a breadth-first search:

import collections
data = [[1,2], [1,8], [5, 8], [3, 5]]
def bfs(d, _start = lambda x:x[0][0]):
  _queue, _result, _seen = collections.deque([_start(d)]), collections.defaultdict(list), []
  while _queue:
    v = _queue.popleft()
    r = [i[0] if i[-1] == v else i[-1] for i in d if v in i]
    _r = [i for i in r if i not in _seen]
    _result[v].extend(_r)
    _queue.extend(_r)
    _seen.extend(_r+[v])
  return _result

print(bfs(data))

Output:

defaultdict(<class 'list'>, {1: [2, 8], 8: [5], 5: [3], 2: [], 3: []})

Upvotes: 0

Dani Mesejo
Dani Mesejo

Reputation: 61910

The most pythonic way I can think, also the fastest one is the following:

def dfs(es, s):
    graph = {}
    for x, y in es:
        graph.setdefault(x, set()).add(y)
        graph.setdefault(y, set()).add(x)

    tree, stack = {}, [s]
    while stack:
        parent = stack.pop()
        children = graph[parent] - tree.keys()
        tree[parent] = list(children)
        stack.extend(children)
    return tree

edges = [[1, 2], [5, 8], [1, 8], [3, 5]]
print(dfs(edges, 1))

Output

{8: [5], 1: [8, 2], 2: [], 3: [], 5: [3]}

The above approach is linear in the size of the graph O(N + E) where N is the number of nodes and E is the number of edges. A simpler approach, albeit slower is:

egdes = [[1, 2], [1, 8], [5, 8], [3, 5]]
tree = {}
sources = {1}


while sources:
    parent = sources.pop()
    children = [t for s, t in egdes if (s == parent and t not in tree)] + \
               [s for s, t in egdes if (t == parent and s not in tree)]
    tree[parent] = children
    sources.update(children)

print(tree)

Output

{8: [5], 1: [2, 8], 2: [], 3: [], 5: [3]}

A faster approach is to delete the already seen edges:

while sources:
    parent = sources.pop()
    children = [y if x == parent else x for x, y in edges if parent in (x, y)]
    edges = [edge for edge in edges if parent not in edge]
    tree[parent] = children
    sources.update(children)

Output

{8: [5], 1: [2, 8], 2: [], 3: [], 5: [3]}

Upvotes: 1

blhsing
blhsing

Reputation: 106488

d = [[1,2], [1,8], [5, 8], [3, 5]]
def t(d, r):
    o = {r: []}
    while d:
        e = d.pop(0)
        if e[1] in o:
            e = e[::-1]
        if e[0] in o:
            o[e[0]].append(e[1])
            o[e[1]] = []
        else:
            d.append(e)
    return o
print(t(d, 1))

This outputs:

{1: [2, 8], 2: [], 8: [5], 5: [3], 3: []}

Upvotes: -1

blhsing
blhsing

Reputation: 106488

A recursive approach:

d = [[1,2], [1,8], [5, 8], [3, 5]]
def t(d, r):
    o = {r: []}
    n = []
    for e in d:
        if r in e:
            if e[1] == r:
                e = e[::-1]
            o[e[0]].append(e[1])
        else:
            n.append(e)
    for i in o[r]:
        o.update(t(n, i))
    return o
print(t(d, 1))

This outputs:

{1: [2, 8], 2: [], 8: [5], 5: [3], 3: []}

Upvotes: 0

user3483203
user3483203

Reputation: 51165

By using a set to keep track of which elements have been added to the dictionary, we can solve this problem in O(n) time. The general idea is to loop through the nodes, check if a node is currently in the graph, and if so, we know that it is the parent when adding to the dictionary.

from collections import defaultdict

class MyTree:
    def __init__(self, data):
        self.data = defaultdict(list)
        self.seen = set()

        for node in data:
            self.add_node(node)

        for el in self.seen:
            if el not in self.data:
                self.data[el] = []

    def add_node(self, el):
        x, y = el

        if y in self.seen:
            self.data[y].append(x)
        else:
            self.data[x].append(y)

        self.seen.update(el)

In action:

arr = [[1,2], [1,8], [5, 8], [3, 5]]

x = MyTree(arr)
print(x.data)

defaultdict(<class 'list'>, {1: [2, 8], 8: [5], 5: [3], 2: [], 3: []})

Upvotes: 2

Related Questions