diff --git a/LinkedList/PartitionList.swift b/LinkedList/PartitionList.swift
index 9625a6b5..7e02df78 100644
--- a/LinkedList/PartitionList.swift
+++ b/LinkedList/PartitionList.swift
@@ -1,8 +1,6 @@
 /**
  * Question Link: https://leetcode.com/problems/partition-list/
- * Primary idea: Tail Insert and merge two lists
- *
- * Note: Swift provides "!==" to determine two objects refer to the same reference
+ * Primary idea: Tail Insert and merge two lists, use dummy to avoid edge case
  * 
  * Time Complexity: O(n), Space Complexity: O(1)
  * Definition for singly-linked list.
@@ -17,23 +15,20 @@
  */
 
 class PartitionList {
-    func partition(head: ListNode?, _ x: Int) -> ListNode? {
-        let prevDummy = ListNode(0)
-        var prev = prevDummy
-        let postDummy = ListNode(0)
-        var post = postDummy
+    func partition(_ head: ListNode?, _ x: Int) -> ListNode? {
+        let prevDummy = ListNode(0), postDummy = ListNode(0)
+        var prev = prevDummy, post = postDummy
         
         var node = head
         
         while node != nil {
             if node!.val < x {
                 prev.next = node
-                prev = node!
+                prev = prev.next!
             } else {
                 post.next = node
-                post = node!
+                post = post.next!
             }
-            
             node = node!.next
         }
         
diff --git a/LinkedList/RemoveNthFromEnd.swift b/LinkedList/RemoveNthFromEnd.swift
index 17ecade3..0b9e7421 100644
--- a/LinkedList/RemoveNthFromEnd.swift
+++ b/LinkedList/RemoveNthFromEnd.swift
@@ -15,11 +15,7 @@
  */
 
 class RemoveNthFromEnd {
-    func removeNthFromEnd(head: ListNode?, _ n: Int) -> ListNode? {
-        guard let head = head else {
-            return nil
-        }
-    
+    func removeNthFromEnd(_ head: ListNode?, _ n: Int) -> ListNode? {
         let dummy = ListNode(0)
         dummy.next = head
         var prev: ListNode? = dummy
@@ -27,15 +23,11 @@ class RemoveNthFromEnd {
         
         // move post
         for _ in 0 ..< n {
-            if post == nil {
-                break
-            }
-            
             post = post!.next
         }
         
         // move prev and post at the same time
-        while post != nil && post!.next != nil {
+        while post!.next != nil {
             prev = prev!.next
             post = post!.next
         }