unitx_interview

question1.py

# A customer of UnitX is trying to identify some defects on a surface. Our neural network is 
# capable of locating the defects on an image and returns a list of points representing
# the location of the defects. E.g. [(102, 215), (109, 201), (50, 32)]

# The customer has an extra requirement where if the distance of two points are smaller than x 
# pixels, then the points should belong to the same group. Write a function that groups the list of points.

# Input:

# points (A list of tuples)
# max_distance (If the distance between two points is smaller than this number, the points are in the same group)
# Output:

# List of list of points. Each sub list represents a group.


# %%
from typing import List, Tuple

from collections import defaultdict

# Follow-up 1
# deal with distributed compute
# subgroup = []points, []points_connected_with_neighbor_subgroup/frontier

# groups = grouped_points(sub_group)
# for g in groups:
#     if intersect(g, points_connected_with_neighbor_subgroup):
#         aggregated_the_group

def v2(p1, p2):
    x1, y1 = p1
    x2, y2 = p2
    return ((x1 - x2) ** 2 + (y1 - y2) **2)**0.5

def grouped_points(points: List[Tuple[int, int]], max_distance: int) -> List[List[Tuple[int, int]]]:
    visited = set() # memory O(V)

    ## O(V^2)
    def neighbors(target):
        candidates = [p for p in points if v2(target, p) <= max_distance]
        return candidates

    ## O(V + E)
    def dfs(p, group):
        visited.add(p) # sync
        group.append(p) 
        for np in neighbors(p):
            if np not in visited:
                dfs(np, group)

    result = []
    for p in points:
        if p not in visited:
            group = []
            dfs(p, group)
            result.append(group)
    return result

if __name__ == "__main__":
    scenarios = [
        {
            'scenario': ([(1, 1)], 1),
            'result': [[(1, 1)]],
        },
        {
            'scenario': ([(1, 1), (0, 0)], 1),
            'result': [[(1, 1)], [(0, 0)]]
        },
        {
            'scenario': ([(1, 1), (0, 0)], 2),
            'result': [[(1, 1), (0, 0)]],
        },
        {
            'scenario': ([(1, 2), (1, 1), (0, 0)], 1),
            'result': [[(1, 1), (1, 2)], [0, 0]],
        },
    ]

    for s in scenarios:
        result = grouped_points(*s['scenario'])
        print(result)
        # assert result == s['result']