61: Re-submit

sankichi92 · Dec 30, 2023 · ed5f462 · ed5f462
1 parent b8c467a
commit ed5f462
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/61-house-robber/src/lib.rs b/61-house-robber/src/lib.rs
@@ -1,4 +1,12 @@
 pub fn rob(nums: Vec<i32>) -> i32 {
+    if nums.len() == 1 {
+        return nums[0];
+    }
+
+    if nums.len() == 2 {
+        return nums[0].max(nums[1]);
+    }
+
     let (mut cand1, mut cand2, mut prev) = (nums[0], nums[1], nums[0] + nums[2]);
 
     for num in nums.iter().skip(3) {
@@ -21,4 +29,9 @@ mod tests {
     fn case2() {
         assert_eq!(rob(vec![2, 7, 9, 3, 1]), 12);
     }
+
+    #[test]
+    fn failed_case1() {
+        assert_eq!(rob(vec![0]), 0);
+    }
 }