用函数解决这个

#include<bits/stdc++.h>
using namespace std;
int digit(int n,int k){
int b,i=1;
while(i<=k){
b=n%10;
n=n/10;
i++;
}
return b;
}
int main() {
int n,k;
cin>>n>>k;
cout<<digit(n,k);
return 0;
}