learning/l33tcode/count-servers-that-communicate.py at master · fspv/learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from typing import List
from collections import defaultdict


class Solution:
    def countServers(self, grid: List[List[int]]) -> int:
        rows, cols = defaultdict(set), defaultdict(set)
        connected = set()

        for row in range(len(grid)):
            for col in range(len(grid[0])):
                if grid[row][col]:
                    if row in rows or col in cols:
                        if len(rows[row]) == 1:
                            connected.add(list(rows[row])[0])
                        if len(cols[col]) == 1:
                            connected.add(list(cols[col])[0])
                        connected.add((row, col))
                    rows[row].add((row, col))
                    cols[col].add((row, col))

        return len(connected)


    def countServersNoMem(self, grid: List[List[int]]) -> int:
        # left-right
        for row in range(len(grid)):
            first = None
            for col in range(len(grid[0])):
                if grid[row][col]:
                    if first is not None:
                        grid[row][col] += 1
                        grid[row][first] += 1
                    else:
                        first = col

        # top-down
        for col in range(len(grid[0])):
            first = None
            for row in range(len(grid)):
                if grid[row][col]:
                    if first is not None:
                        grid[row][col] += 1
                        grid[first][col] += 1
                    else:
                        first = row

        # Count result
        count = 0
        for col in range(len(grid[0])):
            for row in range(len(grid)):
                count += 1 if grid[row][col] > 1 else 0

        return count


class TestSolution:
    def setup(self):
        self.sol = Solution()

    def test_one(self):
        assert self.sol.countServers([[0]]) == 0
        assert self.sol.countServers([[1]]) == 0

    def test_case1(self):
        assert self.sol.countServers([[1,0],[0,1]]) == 0

    def test_case2(self):
        assert self.sol.countServers([[1,0],[1,1]]) == 3

    def test_case3(self):
        assert self.sol.countServers([[1,1,0,0],[0,0,1,0],[0,0,1,0],[0,0,0,1]]) == 4

    def test_case4(self):
        assert self.sol.countServers([[1,1,0,0],[0,0,0,1],[0,0,0,1],[0,0,0,1]]) == 5

    def test_case5(self):
        assert self.sol.countServers(
            [
                [0,0,1,0,1],
                [0,1,0,1,0],
                [0,1,1,1,0],
                [1,0,0,1,1],
                [0,0,1,1,0]
            ]
        ) == 12