learning/l33tcode/k-closest-points-to-origin.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
import unittest


class Solution:
    def kClosest(self, points, K):
        return sorted(points, key=lambda point: point[0] ** 2 + point[1] ** 2)[:K]


class TestSolution(unittest.TestCase):
    def setUp(self):
        self.sol = Solution()

    def test_empty1(self):
        self.assertListEqual(self.sol.kClosest([], 1), [])

    def test_empty2(self):
        self.assertListEqual(self.sol.kClosest([], 0), [])

    def test_one1(self):
        self.assertListEqual(self.sol.kClosest([[1,1]], 1), [[1,1]])

    def test_one2(self):
        self.assertListEqual(self.sol.kClosest([[1,1]], 2), [[1,1]])

    def test_one3(self):
        self.assertListEqual(self.sol.kClosest([[1,1]], 0), [])

    def test_custom1(self):
        self.assertListEqual(self.sol.kClosest([[1,3],[-2,2]], 1), [[-2,2]])

    def test_custom2(self):
        self.assertListEqual(self.sol.kClosest([[3,3],[5,-1],[-2,4]], 2), [[3,3],[-2,4]])


if __name__ == "__main__":
    unittest.main()