diff --git a/src/Align.jsx b/src/Align.jsx
index fba3751..2305652 100644
--- a/src/Align.jsx
+++ b/src/Align.jsx
@@ -143,9 +143,7 @@ class Align extends Component {
         result = alignPoint(source, point, align);
       }
 
-      if (activeElement && activeElement !== document.activeElement) {
-        restoreFocus(activeElement, source);
-      }
+      restoreFocus(activeElement, source);
 
       if (onAlign) {
         onAlign(source, result);
diff --git a/src/util.js b/src/util.js
index 3f51b0b..0045334 100644
--- a/src/util.js
+++ b/src/util.js
@@ -1,3 +1,5 @@
+import contains from 'rc-util/lib/Dom/contains';
+
 export function buffer(fn, ms) {
   let timer;
 
@@ -44,16 +46,12 @@ export function isSimilarValue(val1, val2) {
 }
 
 export function restoreFocus(activeElement, container) {
-  let currentElement = activeElement;
-  while (currentElement) {
-    if (currentElement === container) {
-      break;
-    }
-    currentElement = currentElement.parentNode;
+  if (!activeElement || activeElement === document.activeElement) {
+    return;
   }
 
   // Focus back if is in the container
-  if (currentElement) {
+  if (contains(container, activeElement)) {
     activeElement.focus();
   }
 }