Median of Two Sorted Arrays

 There are two sorted arrays A and B of size m and n respectively. Find the median of the two sorted arrays. The overall run time complexity should be O(log (m+n)).
 1 #include <iostream>
 2 #include <vector>
 3 #include <algorithm>
 4 #include <iterator>
 5 #include <unordered_map>
 6 
 7 using namespace std;
 8 
 9 class Solution {
10 public:
11     double findMedianSortedArrays(int A[], int m, int B[], int n) {
12         int total = m + n;
13         
14         if(total & 0x01)// odd
15             return find_kth(A, m, B, n, total/2+1);
16         else// even
17             return (find_kth(A, m, B, n, total/2) + find_kth(A, m, B, n, total/2+1))/2.0;
18     }
19 private:
20     static double find_kth1(int A[], int m, int B[], int n, int k){
21         int i = 0, j = 0, count = 0;
22         
23         if(0==m) return B[k-1];
24         if(0==n) return A[k-1];
25         
26         while(i < m && j < n){
27             if(A[i] <= B[j]){
28                 ++count;
29                 if(k==count)
30                     return A[i];
31                 ++i;
32             }
33             if(A[i] >= B[j]) {
34                 ++count;
35                 if(k==count)
36                     return B[j];
37                 ++j;
38             }
39         }
40         if(i>m-1){
41             return B[j+k-count-1];
42         }
43         if(j>n-1){
44             return A[i+k-count-1];
45         }
46     }
47     static double find_kth(int A[], int m, int B[], int n, int k){
48         
49         if(0==m) return B[k-1];
50         if(0==n) return A[k-1];
51 
52         int alo = 0, ahi = m - 1;
53         int blo = 0, bhi = n - 1;
54         int count = 0;
55 
56         while(alo <= ahi && blo <= bhi){
57             int amid = alo + (ahi - alo) / 2;
58             int bmid = blo + (bhi - blo) / 2;
59 
60             if(A[amid] <= B[bmid]){
61                 count += amid - alo + 1;
62                 if(k == count)
63                     return A[amid];
64                 alo = amid + 1;
65             }
66             if(A[amid] >= B[bmid]){
67                 count += bmid - blo + 1;
68                 if(k == count)
69                     return B[bmid];
70                 blo = bmid + 1;
71             }
72         }
73         if(alo>m-1){
74             return B[blo+k-count-1];
75         }
76         if(blo>n-1){
77             return A[alo+k-count-1];
78         }
79     }
80 };
81 
82 int main()
83 {
84    Solution s;
85    //int A[] = {1};
86    //int B[] = {2};
87    //int A[] = {1, 2};
88    //int B[] = {2};
89    int A[] = {7};
90    int B[] = {1, 2, 3, 4, 5, 6};
91    double res = s.findMedianSortedArrays(A, 1, B, 6);
92    //double res = s.findMedianSortedArrays(A, 1, B, 1);
93 
94    cout << res;
95    
96    return 0;
97 }