CODEWITHSUNDEEP
pythonlistcolors
Wu Lee
Wu Lee

Reputation: 9

how to return color with correct pair?

Given a list of the following format (id: int, type: "left" | "right", color: string)

return the pair for each left and right id that has the same color.

from collections import defaultdict

def pair(color_list):
    result_dict = defaultdict(lambda: [None, None])
    
    for num, dir, color in color_list:
        if dir == 'left' and result_dict[color][0] is None:
            result_dict[color][0] = num
        elif dir == 'right' and result_dict[color][1] is None:
            result_dict[color][1] = num
            
    result = [x for x in result_dict.values() if None not in x]
    return result

if __name__ == '__main__':
    color_list_1 = [
        (1, "left", "red"),
        (2, "right", "red"),
        (3, "left", "blue"),
        (4, "right", "blue"),
        (5, "left", "yellow")
    ]
    print(pair(color_list_1)) # expected: [[1,2], [3,4]] ; actual: [[1,2], [3,4]]
    
    color_list_2 = [
        (1, "left", "red"),
        (2, "right", "red"),
        (3, "left", "blue"),
        (4, "right", "blue"),
        (5, "left", "blue")
    ]
    print(pair(color_list_2)) # expected: [[1,2], [3,4]] ; actual: [[1, 2], [3, 4]]
    
color_list_3 = [
    (1, "left", "red"),
    (2, "right", "red"),
    (3, "right", "red"),
    (4, "right", "blue"),
    (5, "left", "blue"),
    (6, "right", "blue")
]
print(pair(color_list_3)) # expected: [[1,2], [5,6]] ; actual: [[1, 3], [5, 4]]
    
    color_list_4 = []
    print(pair(color_list_4)) # expected: [] ; actual: []

This code couldn't pass the 3rd case. How should I fix the code to pass 3rd case?

The pair has to be 'left' -> 'right'

Upvotes: 1

Views: 54

Answers (1)

Andrej Kesely
Andrej Kesely

Reputation: 195438

Try:

def state_machine(dct):
    for v in dct.values():
        idx1 = None
        for idx, side in v:
            if side == "left":
                idx1 = idx
            elif side == "right" and not idx1 is None:
                yield [idx1, idx]
                idx1 = None


def pair(color_list):
    tmp = {}
    for idx, side, color in color_list:
        tmp.setdefault(color, []).append((idx, side))

    return list(state_machine(tmp))


if __name__ == "__main__":
    color_list_1 = [
        (1, "left", "red"),
        (2, "right", "red"),
        (3, "left", "blue"),
        (4, "right", "blue"),
        (5, "left", "yellow"),
    ]

    assert pair(color_list_1) == [[1, 2], [3, 4]]

    color_list_2 = [
        (1, "left", "red"),
        (2, "right", "red"),
        (3, "left", "blue"),
        (4, "right", "blue"),
        (5, "left", "blue"),
    ]

    assert pair(color_list_2) == [[1, 2], [3, 4]]

    color_list_3 = [
        (1, "left", "red"),
        (2, "right", "red"),
        (3, "right", "red"),
        (4, "right", "blue"),
        (5, "left", "blue"),
        (6, "right", "blue"),
    ]
    assert pair(color_list_3) == [[1, 2], [5, 6]]

    color_list_4 = []

    assert pair(color_list_4) == []

Result: pass all tests.

EDIT: Updated code according new input.

Upvotes: 1

Related Questions