Java 17 ms O(n + m + k^2 or 3?) worst case merge with lazy computation
I must admit I found this problem the most difficult so far, and I've solved more than 80 already.\n\nAt first my idea was very similar to [this solution][1], w
I must admit I found this problem the most difficult so far, and I've solved more than 80 already.\n\nAt first my idea was very similar to [this solution][1], which I think is the most awesome so far, but I wasn't able to get it right.\n\nSo in the end I had to look up solutions, and I usually don't do that. But once I saw the most common type of solutions so far, I was able to use some of my earlier ideas and came up with this:\n\n    public int[] maxNumber(int[] nums1, int[] nums2, final int k) {\n        if (k > nums1.length + nums2.length)\n            throw new IllegalArgumentException("k=" + k + " n=" + nums1.length + " m=" + nums2.length);\n        int[][] digitPositions1 = findDigitPositions(nums1);\n        int[][] digitPositions2 = findDigitPositions(nums2);\n        Merger best = null;\n        final int l1start = nums2.length >= k ? 0 : k - nums2.length;\n        final int l1end = Math.min(k, nums1.length);\n        // Will now try to get the largest sequences of digits from nums1, nums2 of lengths l1, (k - l1)\n        // and merge them. The largest merge result will be the answer.\n        for (int l1 = l1start; l1 <= l1end; ++l1) {\n            MaxFinder res1 = new MaxFinder(nums1, l1, k, digitPositions1);\n            MaxFinder res2 = new MaxFinder(nums2, k - l1, k, digitPositions2);\n            Merger merger = new Merger(res1, res2, k);\n            if (best == null || merger.compareTo(best) > 0) {\n                best = merger;\n            }\n        }\n        assert best != null;\n        return best.toIntArray();\n    }\n\n    private static int[][] findDigitPositions(int[] nums) {\n        int[][] digitPositions = new int[nums.length][];\n        int[] pos = new int[10];\n        Arrays.fill(pos, -1);\n        for (int i = nums.length - 1; i >= 0; --i) {\n            pos[nums[i]] = i;\n            digitPositions[i] = Arrays.copyOf(pos, pos.length);\n        }\n        return digitPositions;\n    }\n    \n    private static class Merger implements Comparable<Merger> {\n        private final MaxFinder seq1;\n        private final MaxFinder seq2;\n        private final int k; // total length\n        private final int[] output;\n        private int outPos = 0, inPos1 = 0, inPos2 = 0;\n\n        Merger(MaxFinder seq1, MaxFinder seq2, int k) {\n            this.seq1 = seq1;\n            this.seq2 = seq2;\n            this.k = k;\n            this.output = new int[k];\n        }\n        \n        int get(int i) {\n            while (outPos <= i) {\n                if (seq1.compareTo(seq2) >= 0) {\n                    output[outPos++] = seq1.get(inPos1++);\n                } else {\n                    output[outPos++] = seq2.get(inPos2++);\n                }\n            }\n            return output[i];\n        }\n        \n        @Override\n        public int compareTo(Merger that) {\n            for (int i = 0; i < k; ++i) {\n                int v1 = this.get(i), v2 = that.get(i);\n                if (v1 > v2) {\n                    return +1;\n                } else if (v1 < v2) {\n                    return -1;\n                }\n            }\n            return 0;\n        }\n        \n        int[] toIntArray() {\n            get(k - 1); // complete generation\n            return output;\n        }\n    }\n    \n    private static class MaxFinder implements Comparable<MaxFinder> {\n        private final int[] input;\n        private final int k; // total length of two sequences, needed for comparison caching\n        private final int[][] digitPositions; // computed by findDigitPositions\n        private final int[] output;\n        private final int[] lastDiffPos;\n        private int inPos = 0, outPos = 0;\n        private int nextIndex = 0;\n        \n        MaxFinder(int[] nums, int length, int k, int[][] digitPositions) {\n            this.input = nums;\n            this.k = k;\n            this.digitPositions = digitPositions;\n            this.output = new int[length];\n            this.lastDiffPos = new int[2 * k + 1];\n        }\n        \n        int get(int i) {\n            this.nextIndex = i + 1; // If we need to compare with the other sequence, here is where we should start.\n            return compute(i);\n        }\n        \n        // helper for get and compareTo\n        private int compute(int i) {\n            while (outPos <= i) {\n                // Look for the next best digit.\n                for (int d = 9; d >= 0; --d) {\n                    final int nextDigitPos = digitPositions[inPos][d];\n                    // The digit must satisfy two conditions:\n                    // 1. It must be out there somewhere (obviously).\n                    // 2. There must be enough digits after it, or we may end up running out of digits too soon.\n                    if (nextDigitPos != -1 && input.length - (nextDigitPos + 1) >= output.length - outPos - 1) {\n                        output[outPos++] = d;\n                        inPos = digitPositions[inPos][d] + 1;\n                        break;\n                    }\n                }\n            }\n            return output[i];\n        }\n        \n        @Override\n        public int compareTo(MaxFinder that) {\n            final int cacheIndex = k + (this.nextIndex - that.nextIndex);\n            // We can use the cached position if there is one, but only if we haven't advanced past the last found\n            // difference position (which is stored in the cache).\n            if (lastDiffPos[cacheIndex] != 0 && lastDiffPos[cacheIndex] >= this.nextIndex) {\n                int lastDiffPos1 = lastDiffPos[cacheIndex];\n                int lastDiffPos2 = lastDiffPos1 + (that.nextIndex - this.nextIndex);\n                if (lastDiffPos1 < this.output.length && lastDiffPos2 < that.output.length) {\n                    return Integer.compare(this.output[lastDiffPos1], that.output[lastDiffPos2]);\n                } else {\n                    return Integer.compare(this.output.length - this.nextIndex, that.output.length - that.nextIndex);\n                }\n            }\n            // No luck in the cache.\n            final int end = Math.min(this.output.length - this.nextIndex, that.output.length - that.nextIndex);\n            int result = 0;\n            for (int i = 0; i < end; ++i) {\n                result = this.compute(this.nextIndex + i) - that.compute(that.nextIndex + i);\n                if (result != 0) {\n                    lastDiffPos[cacheIndex] = this.nextIndex + i;\n                    break;\n                }\n            }\n            if (result == 0) {\n                lastDiffPos[cacheIndex] = this.nextIndex + end;\n                return Integer.compare(this.output.length - this.nextIndex, that.output.length - that.nextIndex);\n            } else {\n                return result;\n            }\n        }\n    }\n\nThe code is quite verbose, but the ideas are pretty simple:\n\n1. Precompute the table of digit positions so we don't have to search for the next best digit (only have to try from 9 to 0).\n2. Compute everything lazily. We don't know in advance how many digits we need to compare two sequences, so why generate them fully? The same thing for merging.\n3. Cache comparison results. When we've already compared two sequences, we don't want to compare them again if we advance both positions by the same amount. Although some suffix tables would probably be better (maybe compute them lazily too).\n\nThe resulting code runs in 17 ms. The code without improvements 2\u20133 ran for 26 ms, but then I introduced classes and it was 37 ms (because local variables moved to heap?). So these improvements turned 37 to 17, actually. Moving back to local variables will probably improve a lot, but the code would be a real mess.\n\nP. S. I'm not exactly sure: doesn't this comparison caching technique essentially replaces the suffix array and reduces the complexity to O(n + m + k^2). On the 500/500/500 case `MaxFinder::get` is called 575356 times which is of order (500^2) * 2.\n\n  [1]: https://leetcode.com/discuss/75655/strictly-o-nk-c-solution-with-detailed-explanation
['Java']