From 08a12f9a9a9d1e5c9256cd2260332b032c07a50b Mon Sep 17 00:00:00 2001 From: Saffat Hasan Date: Fri, 20 Aug 2021 15:19:36 -0400 Subject: [PATCH] feat(LuckyQuadruplets): impl bfs, wip dp --- .gitignore | 1 + LuckyQuadruplets/bfs.py | 69 ++++++++++++++++++++++++++ LuckyQuadruplets/dp.py | 103 +++++++++++++++++++++++++++++++++++++++ LuckyQuadruplets/main.py | 13 +++++ 4 files changed, 186 insertions(+) create mode 100644 LuckyQuadruplets/bfs.py create mode 100644 LuckyQuadruplets/dp.py create mode 100644 LuckyQuadruplets/main.py diff --git a/.gitignore b/.gitignore index 4c066af..cd2bc23 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ +*.pyc cover.out coverage.html diff --git a/LuckyQuadruplets/bfs.py b/LuckyQuadruplets/bfs.py new file mode 100644 index 0000000..fcbab4d --- /dev/null +++ b/LuckyQuadruplets/bfs.py @@ -0,0 +1,69 @@ +import math + +def LuckyQuadruplets(tree: [int]): + graph = makeGraph(tree) + + # print(graph) + + result = 0 + for i in range(len(tree)): + result += bfs(i, graph) + return result + +def bfs(node: int, graph: {int: [int]}): + visited = set() + visited.add(node) + + # print(f"\nExpanding at node {node}") + + # initially populate queue with nodes that are 1 distance away + # i.e. the neighbors of the current node + queue = [neighbor for neighbor in graph[node]] + result = 0 + distance = 0 + while len(visited) < len(graph) and len(queue) > 0: + distance += 1 + # print(f"{queue=} for {distance=}") + # each layer in BFS represents + # the number of nodes that are equidistant to + # the current node + layer_size = len(queue) + if layer_size >= 3: + # print(f"{layer_size=} adding {quadruplets(layer_size)=} quadruplets...") + result += quadruplets(layer_size) # number of quadruplets we can make with current layer + + # expand + for _ in range(layer_size): + current_node = queue[0] + queue = queue[1:] + visited.add(current_node) + + for neighbor in graph[current_node]: + if neighbor in visited: + continue + queue.append(neighbor) + + return result + +def quadruplets(node_count): + return math.comb(node_count, 3) + +def makeGraph(tree: [int]): + result = {i: [] for i in range(len(tree))} + + for node in range(len(tree)): + parent = getParent(tree, node) + + if parent is None: + continue + result[node].append(parent) + result[parent].append(node) + + return result + +def getParent(tree: [int], node: int) -> int: + if tree[node] == 0: + # root has no parent + return None + + return tree[node] - 1 diff --git a/LuckyQuadruplets/dp.py b/LuckyQuadruplets/dp.py new file mode 100644 index 0000000..8182d87 --- /dev/null +++ b/LuckyQuadruplets/dp.py @@ -0,0 +1,103 @@ +import math + +def LuckyQuadruplets(tree: [int]) -> int: + # assume tree is not null + # assume tree is sorted + total = 0 + children = getChildren(tree) + # for i in range(len(tree)): + # print(f"children[{i}]={children[i]}") + for node in range(len(tree)): + for distance in range(int(math.log2(len(tree) + 1))): + total += expand(tree, children, distance, node) + return total + +def expand(tree: [int], children: {int:[int]}, distance: int, node: int): + parent = getParent(tree, node) + # base case + total = children[node][distance] if distance < len(children[node]) else 0 + + # root does not look "up", only "down" + if parent is None: + return quadruplets(total) + + for offset in range(1, distance): + # e.g. 2 away might look at parent's children excluding itself (node.parent[1] - current[0]) + # or it might look at grndpa's children excluding parent (node.parent.parent[0]) + if distance - offset == 0: + total += 1 + break + children_count = expandHelper(children, distance - offset, node, parent) + if children_count is None: + break + + # Don't update parent pointer if we're at root + if getParent(tree, parent) is None: + continue + node = parent + parent = getParent(tree, node) + + return quadruplets(total) + + +def expandHelper(children: {int:[int]}, distance: int, child: int, parent: int): + parents_children = children[parent] + childs_children = children[child] + + # cannot look down from parent + # because distance is out of range + if distance - 1 >= len(parents_children): + return None + + choices = parents_children[distance-1] + + # remove backtracking / duplicate nodes + # i.e. node->parent->child is itself, is not 2 away + if distance < len(childs_children): + choices -= childs_children[distance] + + return choices + +def quadruplets(number_of_choices): + # fix current node as center + # return all combinations of 3 for the children + return math.comb(number_of_choices, 3) + +def getChildren(tree: [int]) -> {int:[int]}: + children = {} + for i in range(len(tree)-1, 0, -1): + if i in children: + children[i][0] = 1 + else: + children[i] = [1] + + parent = getParent(tree, i) + if parent is None: + continue + + if parent in children: + children[parent] = mergeChildren(children[parent], children[i]) + else: + children[parent] = mergeChildren([1], children[i]) + return children + +def getParent(tree: [int], node: int) -> int: + # root has no parent + if tree[node] == 0: + return None + return tree[node]-1 + +def mergeChildren(parent: [int], child: [int]) -> [int]: + """ + Given two lists of integers, merge the child into parent + with 1 distance of offset (to account for parent being 1 away) + """ + result = [0 for _ in range(max(len(parent), len(child)+1))] + + for i in range(len(parent)): + result[i] = parent[i] + + for i in range(len(child)): + result[i+1] += child[i] + + return result \ No newline at end of file diff --git a/LuckyQuadruplets/main.py b/LuckyQuadruplets/main.py new file mode 100644 index 0000000..dc041ba --- /dev/null +++ b/LuckyQuadruplets/main.py @@ -0,0 +1,13 @@ +import dp, bfs + +def print_result(data): + print(f"Running LuckyQuadruplets against {data=}") + print(f"BFS got {bfs.LuckyQuadruplets(data)}") + print(f"DP got {dp.LuckyQuadruplets(data)}") + +data = [0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9] +current = data[:3] +for i in range(3, len(data)): + current.append(data[i]) + print_result(current) +