Kth element of Two Sorted Arrays

 1 public class FindKthSorted {
 2     //Method 1: time complexity O(k), no extra space
 3
 4     public int findKth (int[] a, int[] b, int k) {
 5         int lenA = a.length;
 6         int lenB = b.length;
 7         if (k > lenA + lenB) {
 8             throw new RuntimeException("Cannot find");
 9         }
10         int indexA = lenA - 1;
11         int indexB = lenB - 1;
12         int count = 0;
13         if (k == 1) {
14             return Math.max(a[indexA], b[indexB]);
15         }
16
17         while (indexA >= 0 && indexB >= 0) {
18             if (a[indexA] < b[indexB]) {
19                 count++;
20                 if (count == k) {
21                     return b[indexB];
22                 }
23                 indexB--;
24             }else{
25                 count++;
26                 if (count == k) {
27                     return a[indexA];
28                 }
29                 indexA--;
30             }
31         }
32         if (indexA < 0) {
33             return b[indexB - (k - count) + 1];
34         }else{
35             return a[indexA - (k - count) + 1];
36         }
37     }
38
39     //Method 2: time complexity (logn + logm)
40     public int findKth2 (int[] a, int[] b, int k) {
41         return helper(a, 0, b, 0, a.length + b.length - k + 1);
42     }
43
44     private int helper(int A[], int startA, int B[], int startB, int k) {
45         //base case:
46         //1. if startA is out of A
47         if (startA >= A.length) {
48             return B[startB - 1 + k];
49         }
50         //2. if startB greater than length
51         if (startB >= B.length) {
52             return A[startA - 1 + k];
53         }
54         //3. if k == 1
55         if (k == 1) {
56             return Math.min(A[startA], B[startB]);
57         }
58         //BUG: pay attention for "-1" for index from 0;
59         int keyA = ((startA + (k / 2) - 1) < A.length) ? A[startA + k / 2 - 1] : Integer.MAX_VALUE;
60         int keyB = ((startB + (k / 2) - 1) < B.length) ? B[startB + k / 2 - 1] : Integer.MAX_VALUE;
61
62         if (keyA > keyB) {
63             return helper(A, startA, B, startB + k / 2, k - k / 2);
64         }else{
65             return helper(A, startA + k / 2, B, startB, k - k / 2);
66         }
67     }
68
69     public static void main(String[] args) {
70         int[] a = {1, 2, 3};
71         int[] b = {4,5};
72         FindKthSorted test = new FindKthSorted();
73         int res = test.findKth2(b, a, 3);
74         System.out.println(res);
75     }
76 }