Federated learning is a distributed computing framework aiming at finding a shared model parameter while protecting the privacy of local agents by sharing only locally updated model parameters without sharing local data with a central server. Through an iterative procedure between agent-side local updates and central server-side aggregation, federated learning reaches its maximum performance after sufficient iterations which is the the possible best performance via central learning. In practice, the statistical heterogeneity of the client’s local dataset (non-iid) degrades the performance and convergence speed of federated learning compared to the iid case. In this paper, we propose two schemes: ( i ) server-driven approach; and ( ii ) client-driven approach, both based on greedy algorithm to tackle the statistical heterogeneity of federated learning. Numerical results validate the advantage of the proposed schemes by improving the performance of state-of-the-art federated learning techniques, e.g., FedAvg, FedProx, FedAdam, and FedDyn with reduced wireless communication cost. Further investigation compares the two proposed schemes to show that the server-driven approach possesses better performance while the client-driven approach has better wireless communication efficiency.