Advanced code time complexity prediction approach using contrastive learning

Shinwoo Park, Joonghyuk Hahn, Elizabeth Orwig, Sang-Ki Ko, Yo-Sub Han

Abstract: It is a crucial task to predict the algorithmic time complexity for estimating the efficiency of a software code. Since the problem is known to be undecidable in theory, there is no 100% accurate tools to solve the problem. Even humans often make mistakes when analyzing the time complexity of code, and this process requires considerable effort and time to thoroughly examine the code. Therefore, we aim to develop an automated method for analyzing code time complexity. We observe that solution codes submitted for coding problems in competitive programming contests tend to have similar time complexities due to constraints such as time limits and functional requirements of the problems. Based on this observation, we propose a contrastive learning-based training strategy that aligns solution codes for the same competitive programming problem. Our training strategy clusters codes with similar time complexities by using both natural language problem descriptions and a single reference code per problem as anchors. This design enables the model to capture core algorithmic features such as loops and recursion more accurately. Experiments in three scenarios – in-dataset, cross-dataset, and cross-language – demonstrate substantial gains on pre-trained code models, consistently surpassing existing methods in both accuracy and generalizability. Our proposed training strategy yields an average 12.54% improvement over cross-entropy-based training, and an 8.01% improvement over data augmentation-based contrastive learning.